forked from phoenix/litellm-mirror
Litellm ruff linting enforcement (#5992)
* ci(config.yml): add a 'check_code_quality' step Addresses https://github.com/BerriAI/litellm/issues/5991 * ci(config.yml): check why circle ci doesn't pick up this test * ci(config.yml): fix to run 'check_code_quality' tests * fix(__init__.py): fix unprotected import * fix(__init__.py): don't remove unused imports * build(ruff.toml): update ruff.toml to ignore unused imports * fix: fix: ruff + pyright - fix linting + type-checking errors * fix: fix linting errors * fix(lago.py): fix module init error * fix: fix linting errors * ci(config.yml): cd into correct dir for checks * fix(proxy_server.py): fix linting error * fix(utils.py): fix bare except causes ruff linting errors * fix: ruff - fix remaining linting errors * fix(clickhouse.py): use standard logging object * fix(__init__.py): fix unprotected import * fix: ruff - fix linting errors * fix: fix linting errors * ci(config.yml): cleanup code qa step (formatting handled in local_testing) * fix(_health_endpoints.py): fix ruff linting errors * ci(config.yml): just use ruff in check_code_quality pipeline for now * build(custom_guardrail.py): include missing file * style(embedding_handler.py): fix ruff check
This commit is contained in:
Krish Dholakia
GitHub
parent
3fc4ae0d65
commit
d57be47b0f
263 changed files with 1687 additions and 3320 deletions
|
|
@ -1,517 +0,0 @@
|
|||
# @pytest.mark.skip(reason="AWS Suspended Account")
|
||||
# # Test the following scenarios:
|
||||
# # 1. Generate a Key, and use it to make a call
|
||||
# # 2. Make a call with invalid key, expect it to fail
|
||||
# # 3. Make a call to a key with invalid model - expect to fail
|
||||
# # 4. Make a call to a key with valid model - expect to pass
|
||||
# # 5. Make a call with key over budget, expect to fail
|
||||
# # 6. Make a streaming chat/completions call with key over budget, expect to fail
|
||||
|
||||
|
||||
# # function to call to generate key - async def new_user(data: NewUserRequest):
|
||||
# # function to validate a request - async def user_auth(request: Request):
|
||||
|
||||
# import sys, os
|
||||
# import traceback
|
||||
# from dotenv import load_dotenv
|
||||
# from fastapi import Request
|
||||
|
||||
# load_dotenv()
|
||||
# import os, io
|
||||
|
||||
# # this file is to test litellm/proxy
|
||||
|
||||
# sys.path.insert(
|
||||
# 0, os.path.abspath("../..")
|
||||
# ) # Adds the parent directory to the system path
|
||||
# import pytest, logging, asyncio
|
||||
# import litellm, asyncio
|
||||
# from litellm.proxy.proxy_server import (
|
||||
# new_user,
|
||||
# user_api_key_auth,
|
||||
# user_update,
|
||||
# generate_key_fn,
|
||||
# )
|
||||
|
||||
# from litellm.proxy._types import NewUserRequest, DynamoDBArgs, GenerateKeyRequest
|
||||
# from litellm.proxy.utils import DBClient, hash_token
|
||||
# from starlette.datastructures import URL
|
||||
|
||||
|
||||
# request_data = {
|
||||
# "model": "azure-gpt-3.5",
|
||||
# "messages": [
|
||||
# {"role": "user", "content": "this is my new test. respond in 50 lines"}
|
||||
# ],
|
||||
# }
|
||||
|
||||
|
||||
# @pytest.fixture
|
||||
# def custom_db_client():
|
||||
# # Assuming DBClient is a class that needs to be instantiated
|
||||
# db_args = {
|
||||
# "ssl_verify": False,
|
||||
# "billing_mode": "PAY_PER_REQUEST",
|
||||
# "region_name": "us-west-2",
|
||||
# }
|
||||
# custom_db_client = DBClient(
|
||||
# custom_db_type="dynamo_db",
|
||||
# custom_db_args=db_args,
|
||||
# )
|
||||
# # Reset litellm.proxy.proxy_server.prisma_client to None
|
||||
# litellm.proxy.proxy_server.prisma_client = None
|
||||
|
||||
# return custom_db_client
|
||||
|
||||
|
||||
# def test_generate_and_call_with_valid_key(custom_db_client):
|
||||
# # 1. Generate a Key, and use it to make a call
|
||||
# setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client)
|
||||
# setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
|
||||
# try:
|
||||
|
||||
# async def test():
|
||||
# request = NewUserRequest()
|
||||
# key = await new_user(request)
|
||||
# print(key)
|
||||
|
||||
# generated_key = key.key
|
||||
# bearer_token = "Bearer " + generated_key
|
||||
|
||||
# request = Request(scope={"type": "http"})
|
||||
# request._url = URL(url="/chat/completions")
|
||||
|
||||
# # use generated key to auth in
|
||||
# result = await user_api_key_auth(request=request, api_key=bearer_token)
|
||||
# print("result from user auth with new key", result)
|
||||
|
||||
# asyncio.run(test())
|
||||
# except Exception as e:
|
||||
# pytest.fail(f"An exception occurred - {str(e)}")
|
||||
|
||||
|
||||
# def test_call_with_invalid_key(custom_db_client):
|
||||
# # 2. Make a call with invalid key, expect it to fail
|
||||
# setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client)
|
||||
# setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
|
||||
# try:
|
||||
|
||||
# async def test():
|
||||
# generated_key = "bad-key"
|
||||
# bearer_token = "Bearer " + generated_key
|
||||
|
||||
# request = Request(scope={"type": "http"}, receive=None)
|
||||
# request._url = URL(url="/chat/completions")
|
||||
|
||||
# # use generated key to auth in
|
||||
# result = await user_api_key_auth(request=request, api_key=bearer_token)
|
||||
# pytest.fail(f"This should have failed!. IT's an invalid key")
|
||||
|
||||
# asyncio.run(test())
|
||||
# except Exception as e:
|
||||
# print("Got Exception", e)
|
||||
# print(e.message)
|
||||
# assert "Authentication Error" in e.message
|
||||
# pass
|
||||
|
||||
|
||||
# def test_call_with_invalid_model(custom_db_client):
|
||||
# # 3. Make a call to a key with an invalid model - expect to fail
|
||||
# from litellm._logging import verbose_proxy_logger
|
||||
# import logging
|
||||
|
||||
# verbose_proxy_logger.setLevel(logging.DEBUG)
|
||||
# setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client)
|
||||
# setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
|
||||
# try:
|
||||
|
||||
# async def test():
|
||||
# request = NewUserRequest(models=["mistral"])
|
||||
# key = await new_user(request)
|
||||
# print(key)
|
||||
|
||||
# generated_key = key.key
|
||||
# bearer_token = "Bearer " + generated_key
|
||||
|
||||
# request = Request(scope={"type": "http"})
|
||||
# request._url = URL(url="/chat/completions")
|
||||
|
||||
# async def return_body():
|
||||
# return b'{"model": "gemini-pro-vision"}'
|
||||
|
||||
# request.body = return_body
|
||||
|
||||
# # use generated key to auth in
|
||||
# result = await user_api_key_auth(request=request, api_key=bearer_token)
|
||||
# pytest.fail(f"This should have failed!. IT's an invalid model")
|
||||
|
||||
# asyncio.run(test())
|
||||
# except Exception as e:
|
||||
# assert (
|
||||
# e.message
|
||||
# == "Authentication Error, API Key not allowed to access model. This token can only access models=['mistral']. Tried to access gemini-pro-vision"
|
||||
# )
|
||||
# pass
|
||||
|
||||
|
||||
# def test_call_with_valid_model(custom_db_client):
|
||||
# # 4. Make a call to a key with a valid model - expect to pass
|
||||
# setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client)
|
||||
# setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
|
||||
# try:
|
||||
|
||||
# async def test():
|
||||
# request = NewUserRequest(models=["mistral"])
|
||||
# key = await new_user(request)
|
||||
# print(key)
|
||||
|
||||
# generated_key = key.key
|
||||
# bearer_token = "Bearer " + generated_key
|
||||
|
||||
# request = Request(scope={"type": "http"})
|
||||
# request._url = URL(url="/chat/completions")
|
||||
|
||||
# async def return_body():
|
||||
# return b'{"model": "mistral"}'
|
||||
|
||||
# request.body = return_body
|
||||
|
||||
# # use generated key to auth in
|
||||
# result = await user_api_key_auth(request=request, api_key=bearer_token)
|
||||
# print("result from user auth with new key", result)
|
||||
|
||||
# asyncio.run(test())
|
||||
# except Exception as e:
|
||||
# pytest.fail(f"An exception occurred - {str(e)}")
|
||||
|
||||
|
||||
# def test_call_with_user_over_budget(custom_db_client):
|
||||
# # 5. Make a call with a key over budget, expect to fail
|
||||
# setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client)
|
||||
# setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
|
||||
# from litellm._logging import verbose_proxy_logger, verbose_logger
|
||||
# import logging
|
||||
|
||||
# litellm.set_verbose = True
|
||||
# verbose_logger.setLevel(logging.DEBUG)
|
||||
# verbose_proxy_logger.setLevel(logging.DEBUG)
|
||||
# try:
|
||||
|
||||
# async def test():
|
||||
# request = NewUserRequest(max_budget=0.00001)
|
||||
# key = await new_user(request)
|
||||
# print(key)
|
||||
|
||||
# generated_key = key.key
|
||||
# user_id = key.user_id
|
||||
# bearer_token = "Bearer " + generated_key
|
||||
|
||||
# request = Request(scope={"type": "http"})
|
||||
# request._url = URL(url="/chat/completions")
|
||||
|
||||
# # use generated key to auth in
|
||||
# result = await user_api_key_auth(request=request, api_key=bearer_token)
|
||||
# print("result from user auth with new key", result)
|
||||
|
||||
# # update spend using track_cost callback, make 2nd request, it should fail
|
||||
# from litellm.proxy.proxy_server import (
|
||||
# _PROXY_track_cost_callback as track_cost_callback,
|
||||
# )
|
||||
# from litellm import ModelResponse, Choices, Message, Usage
|
||||
|
||||
# resp = ModelResponse(
|
||||
# id="chatcmpl-e41836bb-bb8b-4df2-8e70-8f3e160155ac",
|
||||
# choices=[
|
||||
# Choices(
|
||||
# finish_reason=None,
|
||||
# index=0,
|
||||
# message=Message(
|
||||
# content=" Sure! Here is a short poem about the sky:\n\nA canvas of blue, a",
|
||||
# role="assistant",
|
||||
# ),
|
||||
# )
|
||||
# ],
|
||||
# model="gpt-35-turbo", # azure always has model written like this
|
||||
# usage=Usage(prompt_tokens=210, completion_tokens=200, total_tokens=410),
|
||||
# )
|
||||
# await track_cost_callback(
|
||||
# kwargs={
|
||||
# "stream": False,
|
||||
# "litellm_params": {
|
||||
# "metadata": {
|
||||
# "user_api_key": hash_token(generated_key),
|
||||
# "user_api_key_user_id": user_id,
|
||||
# }
|
||||
# },
|
||||
# "response_cost": 0.00002,
|
||||
# },
|
||||
# completion_response=resp,
|
||||
# )
|
||||
# await asyncio.sleep(5)
|
||||
# # use generated key to auth in
|
||||
# result = await user_api_key_auth(request=request, api_key=bearer_token)
|
||||
# print("result from user auth with new key", result)
|
||||
# pytest.fail(f"This should have failed!. They key crossed it's budget")
|
||||
|
||||
# asyncio.run(test())
|
||||
# except Exception as e:
|
||||
# error_detail = e.message
|
||||
# assert "Authentication Error, ExceededBudget:" in error_detail
|
||||
# print(vars(e))
|
||||
|
||||
|
||||
# def test_call_with_user_over_budget_stream(custom_db_client):
|
||||
# # 6. Make a call with a key over budget, expect to fail
|
||||
# setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client)
|
||||
# setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
|
||||
# from litellm._logging import verbose_proxy_logger
|
||||
# import logging
|
||||
|
||||
# litellm.set_verbose = True
|
||||
# verbose_proxy_logger.setLevel(logging.DEBUG)
|
||||
# try:
|
||||
|
||||
# async def test():
|
||||
# request = NewUserRequest(max_budget=0.00001)
|
||||
# key = await new_user(request)
|
||||
# print(key)
|
||||
|
||||
# generated_key = key.key
|
||||
# user_id = key.user_id
|
||||
# bearer_token = "Bearer " + generated_key
|
||||
|
||||
# request = Request(scope={"type": "http"})
|
||||
# request._url = URL(url="/chat/completions")
|
||||
|
||||
# # use generated key to auth in
|
||||
# result = await user_api_key_auth(request=request, api_key=bearer_token)
|
||||
# print("result from user auth with new key", result)
|
||||
|
||||
# # update spend using track_cost callback, make 2nd request, it should fail
|
||||
# from litellm.proxy.proxy_server import (
|
||||
# _PROXY_track_cost_callback as track_cost_callback,
|
||||
# )
|
||||
# from litellm import ModelResponse, Choices, Message, Usage
|
||||
|
||||
# resp = ModelResponse(
|
||||
# id="chatcmpl-e41836bb-bb8b-4df2-8e70-8f3e160155ac",
|
||||
# choices=[
|
||||
# Choices(
|
||||
# finish_reason=None,
|
||||
# index=0,
|
||||
# message=Message(
|
||||
# content=" Sure! Here is a short poem about the sky:\n\nA canvas of blue, a",
|
||||
# role="assistant",
|
||||
# ),
|
||||
# )
|
||||
# ],
|
||||
# model="gpt-35-turbo", # azure always has model written like this
|
||||
# usage=Usage(prompt_tokens=210, completion_tokens=200, total_tokens=410),
|
||||
# )
|
||||
# await track_cost_callback(
|
||||
# kwargs={
|
||||
# "stream": True,
|
||||
# "complete_streaming_response": resp,
|
||||
# "litellm_params": {
|
||||
# "metadata": {
|
||||
# "user_api_key": hash_token(generated_key),
|
||||
# "user_api_key_user_id": user_id,
|
||||
# }
|
||||
# },
|
||||
# "response_cost": 0.00002,
|
||||
# },
|
||||
# completion_response=ModelResponse(),
|
||||
# )
|
||||
# await asyncio.sleep(5)
|
||||
# # use generated key to auth in
|
||||
# result = await user_api_key_auth(request=request, api_key=bearer_token)
|
||||
# print("result from user auth with new key", result)
|
||||
# pytest.fail(f"This should have failed!. They key crossed it's budget")
|
||||
|
||||
# asyncio.run(test())
|
||||
# except Exception as e:
|
||||
# error_detail = e.message
|
||||
# assert "Authentication Error, ExceededBudget:" in error_detail
|
||||
# print(vars(e))
|
||||
|
||||
|
||||
# def test_call_with_user_key_budget(custom_db_client):
|
||||
# # 7. Make a call with a key over budget, expect to fail
|
||||
# setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client)
|
||||
# setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
|
||||
# from litellm._logging import verbose_proxy_logger
|
||||
# import logging
|
||||
|
||||
# verbose_proxy_logger.setLevel(logging.DEBUG)
|
||||
# try:
|
||||
|
||||
# async def test():
|
||||
# request = GenerateKeyRequest(max_budget=0.00001)
|
||||
# key = await generate_key_fn(request)
|
||||
# print(key)
|
||||
|
||||
# generated_key = key.key
|
||||
# user_id = key.user_id
|
||||
# bearer_token = "Bearer " + generated_key
|
||||
|
||||
# request = Request(scope={"type": "http"})
|
||||
# request._url = URL(url="/chat/completions")
|
||||
|
||||
# # use generated key to auth in
|
||||
# result = await user_api_key_auth(request=request, api_key=bearer_token)
|
||||
# print("result from user auth with new key", result)
|
||||
|
||||
# # update spend using track_cost callback, make 2nd request, it should fail
|
||||
# from litellm.proxy.proxy_server import (
|
||||
# _PROXY_track_cost_callback as track_cost_callback,
|
||||
# )
|
||||
# from litellm import ModelResponse, Choices, Message, Usage
|
||||
|
||||
# resp = ModelResponse(
|
||||
# id="chatcmpl-e41836bb-bb8b-4df2-8e70-8f3e160155ac",
|
||||
# choices=[
|
||||
# Choices(
|
||||
# finish_reason=None,
|
||||
# index=0,
|
||||
# message=Message(
|
||||
# content=" Sure! Here is a short poem about the sky:\n\nA canvas of blue, a",
|
||||
# role="assistant",
|
||||
# ),
|
||||
# )
|
||||
# ],
|
||||
# model="gpt-35-turbo", # azure always has model written like this
|
||||
# usage=Usage(prompt_tokens=210, completion_tokens=200, total_tokens=410),
|
||||
# )
|
||||
# await track_cost_callback(
|
||||
# kwargs={
|
||||
# "stream": False,
|
||||
# "litellm_params": {
|
||||
# "metadata": {
|
||||
# "user_api_key": hash_token(generated_key),
|
||||
# "user_api_key_user_id": user_id,
|
||||
# }
|
||||
# },
|
||||
# "response_cost": 0.00002,
|
||||
# },
|
||||
# completion_response=resp,
|
||||
# )
|
||||
# await asyncio.sleep(5)
|
||||
# # use generated key to auth in
|
||||
# result = await user_api_key_auth(request=request, api_key=bearer_token)
|
||||
# print("result from user auth with new key", result)
|
||||
# pytest.fail(f"This should have failed!. They key crossed it's budget")
|
||||
|
||||
# asyncio.run(test())
|
||||
# except Exception as e:
|
||||
# error_detail = e.message
|
||||
# assert "Authentication Error, ExceededTokenBudget:" in error_detail
|
||||
# print(vars(e))
|
||||
|
||||
|
||||
# def test_call_with_key_over_budget_stream(custom_db_client):
|
||||
# # 8. Make a call with a key over budget, expect to fail
|
||||
# setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client)
|
||||
# setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
|
||||
# from litellm._logging import verbose_proxy_logger
|
||||
# import logging
|
||||
|
||||
# litellm.set_verbose = True
|
||||
# verbose_proxy_logger.setLevel(logging.DEBUG)
|
||||
# try:
|
||||
|
||||
# async def test():
|
||||
# request = GenerateKeyRequest(max_budget=0.00001)
|
||||
# key = await generate_key_fn(request)
|
||||
# print(key)
|
||||
|
||||
# generated_key = key.key
|
||||
# user_id = key.user_id
|
||||
# bearer_token = "Bearer " + generated_key
|
||||
|
||||
# request = Request(scope={"type": "http"})
|
||||
# request._url = URL(url="/chat/completions")
|
||||
|
||||
# # use generated key to auth in
|
||||
# result = await user_api_key_auth(request=request, api_key=bearer_token)
|
||||
# print("result from user auth with new key", result)
|
||||
|
||||
# # update spend using track_cost callback, make 2nd request, it should fail
|
||||
# from litellm.proxy.proxy_server import (
|
||||
# _PROXY_track_cost_callback as track_cost_callback,
|
||||
# )
|
||||
# from litellm import ModelResponse, Choices, Message, Usage
|
||||
|
||||
# resp = ModelResponse(
|
||||
# id="chatcmpl-e41836bb-bb8b-4df2-8e70-8f3e160155ac",
|
||||
# choices=[
|
||||
# Choices(
|
||||
# finish_reason=None,
|
||||
# index=0,
|
||||
# message=Message(
|
||||
# content=" Sure! Here is a short poem about the sky:\n\nA canvas of blue, a",
|
||||
# role="assistant",
|
||||
# ),
|
||||
# )
|
||||
# ],
|
||||
# model="gpt-35-turbo", # azure always has model written like this
|
||||
# usage=Usage(prompt_tokens=210, completion_tokens=200, total_tokens=410),
|
||||
# )
|
||||
# await track_cost_callback(
|
||||
# kwargs={
|
||||
# "stream": True,
|
||||
# "complete_streaming_response": resp,
|
||||
# "litellm_params": {
|
||||
# "metadata": {
|
||||
# "user_api_key": hash_token(generated_key),
|
||||
# "user_api_key_user_id": user_id,
|
||||
# }
|
||||
# },
|
||||
# "response_cost": 0.00002,
|
||||
# },
|
||||
# completion_response=ModelResponse(),
|
||||
# )
|
||||
# await asyncio.sleep(5)
|
||||
# # use generated key to auth in
|
||||
# result = await user_api_key_auth(request=request, api_key=bearer_token)
|
||||
# print("result from user auth with new key", result)
|
||||
# pytest.fail(f"This should have failed!. They key crossed it's budget")
|
||||
|
||||
# asyncio.run(test())
|
||||
# except Exception as e:
|
||||
# error_detail = e.message
|
||||
# assert "Authentication Error, ExceededTokenBudget:" in error_detail
|
||||
# print(vars(e))
|
||||
|
||||
|
||||
# def test_dynamo_db_migration(custom_db_client):
|
||||
# # Tests the temporary patch we have in place
|
||||
# setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client)
|
||||
# setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
|
||||
# setattr(litellm.proxy.proxy_server, "user_custom_auth", None)
|
||||
# try:
|
||||
|
||||
# async def test():
|
||||
# request = GenerateKeyRequest(max_budget=1)
|
||||
# key = await generate_key_fn(request)
|
||||
# print(key)
|
||||
|
||||
# generated_key = key.key
|
||||
# bearer_token = (
|
||||
# "Bearer " + generated_key
|
||||
# ) # this works with ishaan's db, it's a never expiring key
|
||||
|
||||
# request = Request(scope={"type": "http"})
|
||||
# request._url = URL(url="/chat/completions")
|
||||
|
||||
# async def return_body():
|
||||
# return b'{"model": "azure-models"}'
|
||||
|
||||
# request.body = return_body
|
||||
|
||||
# # use generated key to auth in
|
||||
# result = await user_api_key_auth(request=request, api_key=bearer_token)
|
||||
# print("result from user auth with new key", result)
|
||||
|
||||
# asyncio.run(test())
|
||||
# except Exception as e:
|
||||
# pytest.fail(f"An exception occurred - {traceback.format_exc()}")
|
||||
Loading…
Add table
Add a link
Reference in a new issue