From 5703da9b4253ade7b1d0177f34f640e0ad5d5276 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 8 Aug 2024 10:30:15 -0700 Subject: [PATCH] fix(user_api_key_auth.py): Fixes https://github.com/BerriAI/litellm/issues/5111 --- litellm/proxy/_experimental/out/404.html | 1 - .../proxy/_experimental/out/model_hub.html | 1 - .../proxy/_experimental/out/onboarding.html | 1 - litellm/proxy/_new_secret_config.yaml | 11 --------- litellm/proxy/auth/auth_checks.py | 2 +- litellm/proxy/auth/user_api_key_auth.py | 24 ++++++++++++------- tests/test_keys.py | 20 ++++++++++++++++ 7 files changed, 37 insertions(+), 23 deletions(-) delete mode 100644 litellm/proxy/_experimental/out/404.html delete mode 100644 litellm/proxy/_experimental/out/model_hub.html delete mode 100644 litellm/proxy/_experimental/out/onboarding.html diff --git a/litellm/proxy/_experimental/out/404.html b/litellm/proxy/_experimental/out/404.html deleted file mode 100644 index 4428091a3..000000000 --- a/litellm/proxy/_experimental/out/404.html +++ /dev/null @@ -1 +0,0 @@ -404: This page could not be found.LiteLLM Dashboard

404

This page could not be found.

\ No newline at end of file diff --git a/litellm/proxy/_experimental/out/model_hub.html b/litellm/proxy/_experimental/out/model_hub.html deleted file mode 100644 index 700debc2f..000000000 --- a/litellm/proxy/_experimental/out/model_hub.html +++ /dev/null @@ -1 +0,0 @@ -LiteLLM Dashboard \ No newline at end of file diff --git a/litellm/proxy/_experimental/out/onboarding.html b/litellm/proxy/_experimental/out/onboarding.html deleted file mode 100644 index 16f3f2593..000000000 --- a/litellm/proxy/_experimental/out/onboarding.html +++ /dev/null @@ -1 +0,0 @@ -LiteLLM Dashboard \ No newline at end of file diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 35ef59c96..97457ef84 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -1,15 +1,4 @@ model_list: - - model_name: "gpt-3.5-turbo" - litellm_params: - model: "gpt-3.5-turbo" - model_name: "gpt-4" litellm_params: model: "gpt-4" - api_key: "bad_key" - - model_name: "gpt-4o" - litellm_params: - model: "gpt-4o" - -litellm_settings: - enable_json_schema_validation: true - fallbacks: [{"gpt-3.5-turbo": ["gpt-4", "gpt-4o"]}] diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 180ed6a6e..366eb1fb2 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -375,7 +375,7 @@ async def get_user_object( await user_api_key_cache.async_set_cache(key=user_id, value=_response) return _response - except Exception as e: # if user not in db + except Exception: # if user not in db raise ValueError( f"User doesn't exist in db. 'user_id'={user_id}. Create user via `/user/new` call." ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index ca1a1a787..743a4cd90 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -652,14 +652,22 @@ async def user_api_key_auth( # Check 2. If user_id for this token is in budget - done in common_checks() if valid_token.user_id is not None: - user_obj = await get_user_object( - user_id=valid_token.user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - user_id_upsert=False, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) + try: + user_obj = await get_user_object( + user_id=valid_token.user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception as e: + verbose_logger.warning( + "litellm.proxy.auth.user_api_key_auth.py::user_api_key_auth() - Unable to get user from db/cache. Setting user_obj to None. Exception received - {}".format( + str(e) + ) + ) + user_obj = None # Check 3. Check if user is in their team budget if valid_token.team_member_spend is not None: diff --git a/tests/test_keys.py b/tests/test_keys.py index 877c239bc..322c53192 100644 --- a/tests/test_keys.py +++ b/tests/test_keys.py @@ -798,3 +798,23 @@ async def test_key_model_list(model_access, model_access_level, model_endpoint): elif model_endpoint == "/model/info": assert isinstance(model_list["data"], list) assert len(model_list["data"]) == 1 + + +@pytest.mark.asyncio +async def test_key_user_not_in_db(): + """ + - Create a key with unique user-id (not in db) + - Check if key can make `/chat/completion` call + """ + my_unique_user = str(uuid.uuid4()) + async with aiohttp.ClientSession() as session: + key_gen = await generate_key( + session=session, + i=0, + user_id=my_unique_user, + ) + key = key_gen["key"] + try: + await chat_completion(session=session, key=key) + except Exception as e: + pytest.fail(f"Expected this call to work - {str(e)}")