From edef3cfc3a235f168eca4265df8fdc57c7bc2c16 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Wed, 17 Jan 2024 11:20:07 -0800 Subject: [PATCH 01/11] =?UTF-8?q?bump:=20version=201.17.14=20=E2=86=92=201?= =?UTF-8?q?.17.15?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pyproject.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 6c9b4c69b..5a7947539 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "litellm" -version = "1.17.14" +version = "1.17.15" description = "Library to easily interface with LLM API providers" authors = ["BerriAI"] license = "MIT License" @@ -61,7 +61,7 @@ requires = ["poetry-core", "wheel"] build-backend = "poetry.core.masonry.api" [tool.commitizen] -version = "1.17.14" +version = "1.17.15" version_files = [ "pyproject.toml:^version" ] From acbbccdef7d33c31935087e00197d3c5cd2551ec Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Wed, 17 Jan 2024 11:28:59 -0800 Subject: [PATCH 02/11] v0 testing --- litellm/tests/test_key_generate_prisma.py | 304 ++++++++++++++++++++++ 1 file changed, 304 insertions(+) create mode 100644 litellm/tests/test_key_generate_prisma.py diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py new file mode 100644 index 000000000..e0b0ab2da --- /dev/null +++ b/litellm/tests/test_key_generate_prisma.py @@ -0,0 +1,304 @@ +# 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 +from litellm.proxy.utils import PrismaClient, ProxyLogging + +from litellm.proxy._types import NewUserRequest, DynamoDBArgs +from litellm.proxy.utils import DBClient +from starlette.datastructures import URL +from litellm.caching import DualCache + +proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + +prisma_client = PrismaClient( + database_url=os.environ["PROXY_DATABASE_URL"], proxy_logging_obj=proxy_logging_obj +) + + +request_data = { + "model": "azure-gpt-3.5", + "messages": [ + {"role": "user", "content": "this is my new test. respond in 50 lines"} + ], +} + + +def test_generate_and_call_with_valid_key(): + # 1. Generate a Key, and use it to make a call + prisma_client.connect() + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_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(): + # 2. Make a call with invalid key, expect it to fail + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_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.detail) + assert "Authentication Error" in e.detail + pass + + +def test_call_with_invalid_model(): + # 3. Make a call to a key with an invalid model - expect to fail + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_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.detail + == "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(): + # 4. Make a call to a key with a valid model - expect to pass + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_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_key_over_budget(): + # 5. Make a call with a key over budget, expect to fail + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") + 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 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": generated_key, + "user_api_key_user_id": user_id, + } + }, + }, + completion_response=resp, + ) + + # 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.detail + assert "Authentication Error, ExceededBudget:" in error_detail + print(vars(e)) + + +def test_call_with_key_over_budget_stream(): + # 6. Make a call with a key over budget, expect to fail + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_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 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": generated_key, + "user_api_key_user_id": user_id, + } + }, + }, + completion_response=ModelResponse(), + ) + + # 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.detail + assert "Authentication Error, ExceededBudget:" in error_detail + print(vars(e)) From ec68c3a29f2d9aec1af72fcbdbfba87dc1c7e35f Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Wed, 17 Jan 2024 11:57:48 -0800 Subject: [PATCH 03/11] (feat) improve Prisma Auth Error --- litellm/proxy/utils.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 80c05e78d..c4c76846a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -7,6 +7,7 @@ from litellm.proxy.hooks.parallel_request_limiter import MaxParallelRequestsHand from litellm.proxy.hooks.max_budget_limiter import MaxBudgetLimiter from litellm.integrations.custom_logger import CustomLogger from litellm.proxy.db.base_client import CustomDB +from litellm._logging import verbose_proxy_logger from fastapi import HTTPException, status import smtplib from email.mime.text import MIMEText @@ -381,7 +382,7 @@ class PrismaClient: # Token does not exist. raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, - detail="invalid user key", + detail="Authentication Error: invalid user key - token does not exist", ) elif user_id is not None: response = await self.db.litellm_usertable.find_unique( # type: ignore @@ -559,7 +560,13 @@ class PrismaClient: ) async def connect(self): try: + verbose_proxy_logger.debug( + "PrismaClient: connect() called Attempting to Connect to DB" + ) if self.db.is_connected() == False: + verbose_proxy_logger.debug( + "PrismaClient: DB not connected, Attempting to Connect to DB" + ) await self.db.connect() except Exception as e: asyncio.create_task( From 4036c64bdf4a9620ef0750ed6a8827f7a59c9940 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Wed, 17 Jan 2024 11:58:14 -0800 Subject: [PATCH 04/11] (test) prisma key gen --- litellm/tests/test_key_generate_prisma.py | 21 ++++++++++++++++----- 1 file changed, 16 insertions(+), 5 deletions(-) diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index e0b0ab2da..534a0ca81 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -27,6 +27,9 @@ import pytest, logging, asyncio import litellm, asyncio from litellm.proxy.proxy_server import new_user, user_api_key_auth, user_update from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm._logging import verbose_proxy_logger + +verbose_proxy_logger.setLevel(level=logging.DEBUG) from litellm.proxy._types import NewUserRequest, DynamoDBArgs from litellm.proxy.utils import DBClient @@ -35,10 +38,6 @@ from litellm.caching import DualCache proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) -prisma_client = PrismaClient( - database_url=os.environ["PROXY_DATABASE_URL"], proxy_logging_obj=proxy_logging_obj -) - request_data = { "model": "azure-gpt-3.5", @@ -46,16 +45,22 @@ request_data = { {"role": "user", "content": "this is my new test. respond in 50 lines"} ], } +prisma_client = PrismaClient( + database_url=os.environ["DATABASE_URL"], proxy_logging_obj=proxy_logging_obj +) def test_generate_and_call_with_valid_key(): # 1. Generate a Key, and use it to make a call - prisma_client.connect() + + print("prisma client=", prisma_client) + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") try: async def test(): + await litellm.litellm.proxy.proxy_server.prisma_client.connect() request = NewUserRequest() key = await new_user(request) print(key) @@ -82,6 +87,7 @@ def test_call_with_invalid_key(): try: async def test(): + await litellm.litellm.proxy.proxy_server.prisma_client.connect() generated_key = "bad-key" bearer_token = "Bearer " + generated_key @@ -90,6 +96,7 @@ def test_call_with_invalid_key(): # use generated key to auth in result = await user_api_key_auth(request=request, api_key=bearer_token) + print("got result", result) pytest.fail(f"This should have failed!. IT's an invalid key") asyncio.run(test()) @@ -107,6 +114,7 @@ def test_call_with_invalid_model(): try: async def test(): + await litellm.litellm.proxy.proxy_server.prisma_client.connect() request = NewUserRequest(models=["mistral"]) key = await new_user(request) print(key) @@ -142,6 +150,7 @@ def test_call_with_valid_model(): try: async def test(): + await litellm.litellm.proxy.proxy_server.prisma_client.connect() request = NewUserRequest(models=["mistral"]) key = await new_user(request) print(key) @@ -173,6 +182,7 @@ def test_call_with_key_over_budget(): try: async def test(): + await litellm.litellm.proxy.proxy_server.prisma_client.connect() request = NewUserRequest(max_budget=0.00001) key = await new_user(request) print(key) @@ -244,6 +254,7 @@ def test_call_with_key_over_budget_stream(): try: async def test(): + await litellm.litellm.proxy.proxy_server.prisma_client.connect() request = NewUserRequest(max_budget=0.00001) key = await new_user(request) print(key) From 2273f8055c8677aefcbbd2d7ab6babbc616bce8c Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Wed, 17 Jan 2024 12:23:04 -0800 Subject: [PATCH 05/11] (test) add test fixture for dynamodb testing --- litellm/tests/test_key_generate_dynamodb.py | 40 ++++++++++++--------- 1 file changed, 24 insertions(+), 16 deletions(-) diff --git a/litellm/tests/test_key_generate_dynamodb.py b/litellm/tests/test_key_generate_dynamodb.py index 207c8835a..09f699af7 100644 --- a/litellm/tests/test_key_generate_dynamodb.py +++ b/litellm/tests/test_key_generate_dynamodb.py @@ -32,16 +32,6 @@ from litellm.proxy.utils import DBClient from starlette.datastructures import URL -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, -) - request_data = { "model": "azure-gpt-3.5", "messages": [ @@ -50,7 +40,25 @@ request_data = { } -def test_generate_and_call_with_valid_key(): +@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") @@ -76,7 +84,7 @@ def test_generate_and_call_with_valid_key(): pytest.fail(f"An exception occurred - {str(e)}") -def test_call_with_invalid_key(): +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") @@ -101,7 +109,7 @@ def test_call_with_invalid_key(): pass -def test_call_with_invalid_model(): +def test_call_with_invalid_model(custom_db_client): # 3. Make a call to a key with an invalid model - expect to fail setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client) setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") @@ -136,7 +144,7 @@ def test_call_with_invalid_model(): pass -def test_call_with_valid_model(): +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") @@ -167,7 +175,7 @@ def test_call_with_valid_model(): pytest.fail(f"An exception occurred - {str(e)}") -def test_call_with_key_over_budget(): +def test_call_with_key_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") @@ -233,7 +241,7 @@ def test_call_with_key_over_budget(): print(vars(e)) -def test_call_with_key_over_budget_stream(): +def test_call_with_key_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 62963146fed633e19100b194c4a3c3d16af2ffee Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Wed, 17 Jan 2024 12:23:45 -0800 Subject: [PATCH 06/11] (test) prisma /key/gen --- litellm/tests/test_key_generate_prisma.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 534a0ca81..5b98a886b 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -60,7 +60,7 @@ def test_generate_and_call_with_valid_key(): try: async def test(): - await litellm.litellm.proxy.proxy_server.prisma_client.connect() + await litellm.proxy.proxy_server.prisma_client.connect() request = NewUserRequest() key = await new_user(request) print(key) @@ -87,7 +87,7 @@ def test_call_with_invalid_key(): try: async def test(): - await litellm.litellm.proxy.proxy_server.prisma_client.connect() + await litellm.proxy.proxy_server.prisma_client.connect() generated_key = "bad-key" bearer_token = "Bearer " + generated_key @@ -114,7 +114,7 @@ def test_call_with_invalid_model(): try: async def test(): - await litellm.litellm.proxy.proxy_server.prisma_client.connect() + await litellm.proxy.proxy_server.prisma_client.connect() request = NewUserRequest(models=["mistral"]) key = await new_user(request) print(key) @@ -150,7 +150,7 @@ def test_call_with_valid_model(): try: async def test(): - await litellm.litellm.proxy.proxy_server.prisma_client.connect() + await litellm.proxy.proxy_server.prisma_client.connect() request = NewUserRequest(models=["mistral"]) key = await new_user(request) print(key) @@ -182,7 +182,7 @@ def test_call_with_key_over_budget(): try: async def test(): - await litellm.litellm.proxy.proxy_server.prisma_client.connect() + await litellm.proxy.proxy_server.prisma_client.connect() request = NewUserRequest(max_budget=0.00001) key = await new_user(request) print(key) @@ -254,7 +254,7 @@ def test_call_with_key_over_budget_stream(): try: async def test(): - await litellm.litellm.proxy.proxy_server.prisma_client.connect() + await litellm.proxy.proxy_server.prisma_client.connect() request = NewUserRequest(max_budget=0.00001) key = await new_user(request) print(key) From f1ae6957ca6188aba26ec3a572b88830e6af4383 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Wed, 17 Jan 2024 12:44:06 -0800 Subject: [PATCH 07/11] (test) use fixture for prisma --- litellm/tests/test_key_generate_prisma.py | 28 +++++++++++++++-------- 1 file changed, 19 insertions(+), 9 deletions(-) diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 5b98a886b..3dd5aa709 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -45,12 +45,22 @@ request_data = { {"role": "user", "content": "this is my new test. respond in 50 lines"} ], } -prisma_client = PrismaClient( - database_url=os.environ["DATABASE_URL"], proxy_logging_obj=proxy_logging_obj -) -def test_generate_and_call_with_valid_key(): +@pytest.fixture +def prisma_client(): + # Assuming DBClient is a class that needs to be instantiated + prisma_client = PrismaClient( + database_url=os.environ["DATABASE_URL"], proxy_logging_obj=proxy_logging_obj + ) + + # Reset litellm.proxy.proxy_server.prisma_client to None + litellm.proxy.proxy_server.custom_db_client = None + + return prisma_client + + +def test_generate_and_call_with_valid_key(prisma_client): # 1. Generate a Key, and use it to make a call print("prisma client=", prisma_client) @@ -80,7 +90,7 @@ def test_generate_and_call_with_valid_key(): pytest.fail(f"An exception occurred - {str(e)}") -def test_call_with_invalid_key(): +def test_call_with_invalid_key(prisma_client): # 2. Make a call with invalid key, expect it to fail setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") @@ -107,7 +117,7 @@ def test_call_with_invalid_key(): pass -def test_call_with_invalid_model(): +def test_call_with_invalid_model(prisma_client): # 3. Make a call to a key with an invalid model - expect to fail setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") @@ -143,7 +153,7 @@ def test_call_with_invalid_model(): pass -def test_call_with_valid_model(): +def test_call_with_valid_model(prisma_client): # 4. Make a call to a key with a valid model - expect to pass setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") @@ -175,7 +185,7 @@ def test_call_with_valid_model(): pytest.fail(f"An exception occurred - {str(e)}") -def test_call_with_key_over_budget(): +def test_call_with_key_over_budget(prisma_client): # 5. Make a call with a key over budget, expect to fail setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") @@ -242,7 +252,7 @@ def test_call_with_key_over_budget(): print(vars(e)) -def test_call_with_key_over_budget_stream(): +def test_call_with_key_over_budget_stream(prisma_client): # 6. Make a call with a key over budget, expect to fail setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") From 399f0ba620c0c3b148ce4ca9e2484f7df11aaec4 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Wed, 17 Jan 2024 12:56:26 -0800 Subject: [PATCH 08/11] (fix) prisma - non expiring keys --- litellm/proxy/utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index c4c76846a..ae4aebad5 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -12,6 +12,7 @@ from fastapi import HTTPException, status import smtplib from email.mime.text import MIMEText from email.mime.multipart import MIMEMultipart +from datetime import datetime def print_verbose(print_statement): @@ -376,7 +377,8 @@ class PrismaClient: print_verbose(f"PrismaClient: response={response}") if response is not None: # for prisma we need to cast the expires time to str - response.expires = response.expires.isoformat() + if isinstance(response.expires, datetime): + response.expires = response.expires.isoformat() return response else: # Token does not exist. From 5f7a89d2e7eed07f7f8d0ad1a836baaa88399e60 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Wed, 17 Jan 2024 13:05:23 -0800 Subject: [PATCH 09/11] (test) prisma - non expiring keys --- litellm/tests/test_key_generate_prisma.py | 32 +++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 3dd5aa709..78fbc7037 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -5,6 +5,8 @@ # 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 +# 7. Make a call with an key that never expires, expect to pass +# 8. Make a call with an expired key, expect to fail # function to call to generate key - async def new_user(data: NewUserRequest): @@ -323,3 +325,33 @@ def test_call_with_key_over_budget_stream(prisma_client): error_detail = e.detail assert "Authentication Error, ExceededBudget:" in error_detail print(vars(e)) + + +def test_generate_and_call_with_valid_key_never_expires(prisma_client): + # 7. Make a call with an key that never expires, expect to pass + + print("prisma client=", prisma_client) + + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") + try: + + async def test(): + await litellm.proxy.proxy_server.prisma_client.connect() + request = NewUserRequest(duration=None) + 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)}") From a0eec51ee64bbeceb269723825918c24d8cc85a4 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Wed, 17 Jan 2024 13:24:15 -0800 Subject: [PATCH 10/11] (test) expired key prisma --- litellm/proxy/proxy_server.py | 2 +- litellm/tests/test_key_generate_prisma.py | 34 +++++++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0503d2aad..471e26d3d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -370,7 +370,7 @@ async def user_api_key_auth( # Token exists but is expired. raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, - detail="expired user key", + detail=f"Authentication Error - Expired Key. Key Expiry time {expiry_time} and current time {current_time}", ) # Token passed all checks diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 78fbc7037..7a9e2bddc 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -355,3 +355,37 @@ def test_generate_and_call_with_valid_key_never_expires(prisma_client): asyncio.run(test()) except Exception as e: pytest.fail(f"An exception occurred - {str(e)}") + + +def test_generate_and_call_with_expired_key(prisma_client): + # 8. Make a call with an expired key, expect to fail + + print("prisma client=", prisma_client) + + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") + try: + + async def test(): + await litellm.proxy.proxy_server.prisma_client.connect() + request = NewUserRequest(duration="0s") + 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) + pytest.fail(f"This should have failed!. IT's an expired key") + + asyncio.run(test()) + except Exception as e: + print("Got Exception", e) + print(e.detail) + assert "Authentication Error" in e.detail + pass From cb9a506e68730259b99e7d409169d88137559334 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Wed, 17 Jan 2024 13:38:02 -0800 Subject: [PATCH 11/11] (chore) undo change to pyproject.toml --- pyproject.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 5a7947539..6c9b4c69b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "litellm" -version = "1.17.15" +version = "1.17.14" description = "Library to easily interface with LLM API providers" authors = ["BerriAI"] license = "MIT License" @@ -61,7 +61,7 @@ requires = ["poetry-core", "wheel"] build-backend = "poetry.core.masonry.api" [tool.commitizen] -version = "1.17.15" +version = "1.17.14" version_files = [ "pyproject.toml:^version" ]