diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 0e9edc809e7..c5119924f5b 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -26,8 +26,6 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers -from litellm.proxy.proxy_server import prisma_client, user_api_key_cache -from litellm.proxy.auth.auth_checks import get_user_object # Cache special headers as a frozenset for O(1) lookup performance _SPECIAL_HEADERS_CACHE = frozenset( @@ -739,15 +737,20 @@ class LiteLLMProxyRequestSetup: return user_api_key_dict # Quick email-ish heuristic - if isinstance(header_value, str) and re.match(r"^[^@\s]+@[^@\s]+\.[^@\s]+$", header_value): + if isinstance(header_value, str) and re.match( + r"^[^@\s]+@[^@\s]+\.[^@\s]+$", header_value + ): try: - # Import at runtime to avoid circular imports - if prisma_client is not None: + # Import at runtime to avoid circular imports. + from litellm.proxy import proxy_server + from litellm.proxy.auth.auth_checks import get_user_object + + if proxy_server.prisma_client is not None: try: user_obj = await get_user_object( user_id=str(header_value), - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, + prisma_client=proxy_server.prisma_client, + user_api_key_cache=proxy_server.user_api_key_cache, user_id_upsert=True, user_email=str(header_value), ) @@ -757,7 +760,7 @@ class LiteLLMProxyRequestSetup: if user_obj is not None: # Ensure role is internal_user (best-effort) try: - await prisma_client.db.litellm_usertable.update( + await proxy_server.prisma_client.db.litellm_usertable.update( where={"user_id": user_obj.user_id}, data={"user_role": str(LitellmUserRoles.INTERNAL_USER)}, ) diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index f4d77a64dba..52d9e063ce1 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -2565,7 +2565,8 @@ async def test_add_internal_user_from_user_mapping_no_header_or_mapping_returns_ async def test_add_internal_user_from_user_mapping_resolves_email_header_to_internal_user( monkeypatch, ): - from litellm import proxy as proxy_pkg + import litellm.proxy.auth.auth_checks as auth_checks + import litellm.proxy.proxy_server as proxy_server user_api_key_dict = UserAPIKeyAuth(api_key="test-key") headers = {"X-OpenWebUI-User-Email": "internal@example.com"} @@ -2585,22 +2586,10 @@ async def test_add_internal_user_from_user_mapping_resolves_email_header_to_inte fake_user.user_email = "internal@example.com" fake_user.user_role = "internal_user" - monkeypatch.setattr( - proxy_pkg.litellm_pre_call_utils, - "prisma_client", - fake_prisma_client, - ) - monkeypatch.setattr( - proxy_pkg.litellm_pre_call_utils, - "user_api_key_cache", - fake_cache, - ) + monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma_client) + monkeypatch.setattr(proxy_server, "user_api_key_cache", fake_cache) get_user_object_mock = AsyncMock(return_value=fake_user) - monkeypatch.setattr( - proxy_pkg.litellm_pre_call_utils, - "get_user_object", - get_user_object_mock, - ) + monkeypatch.setattr(auth_checks, "get_user_object", get_user_object_mock) result = await LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping( general_settings, user_api_key_dict, headers