Merge pull request #5117 from BerriAI/litellm_auth_fix

fix(user_api_key_auth.py): Fix issue with key auth w/ user not in db

litellm/proxy/_new_secret_config.yaml

@@ -6,4 +6,3 @@ model_list:
 litellm_settings:
   max_internal_user_budget: 0.001
   internal_user_budget_duration: "5m"
-

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."
         )

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:

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)}")