From 36e5ed2048b9183ab85502ecadf568496da21f36 Mon Sep 17 00:00:00 2001
From: Yassin Kortam <yassin@berri.ai>
Date: Mon, 8 Jun 2026 18:46:38 -0700
Subject: [PATCH] refactor(identity): hoist function-local imports to module
 top

Function-local imports in the new identity module were cargo-culted to
dodge a circular import that does not exist. Move them to module-level.
The one exception is the proxy_server import inside get_identity_cache
which IS a real cycle (proxy_server -> user_api_key_auth -> identity)
and stays lazy.

Test patches that pinned the imported name in the source module move to
the binding in the consuming module, which is correct for module-level
imports anyway.
---
 litellm/identity/adapter.py                  |  8 +------
 litellm/identity/cache.py                    |  5 ++--
 litellm/identity/extractors/api_key.py       |  3 +--
 litellm/identity/extractors/client.py        |  3 +--
 litellm/identity/extractors/end_user.py      |  5 ++--
 litellm/identity/jwt.py                      |  7 ++----
 litellm/identity/oauth2.py                   |  7 ++----
 litellm/identity/resolver.py                 |  3 +--
 litellm/identity/store.py                    | 25 ++++++++++----------
 tests/test_litellm/identity/test_resolver.py |  2 +-
 tests/test_litellm/identity/test_store.py    | 10 ++++----
 11 files changed, 30 insertions(+), 48 deletions(-)

diff --git a/litellm/identity/adapter.py b/litellm/identity/adapter.py
index 39b8a0ffc5e..cd105d70d48 100644
--- a/litellm/identity/adapter.py
+++ b/litellm/identity/adapter.py
@@ -7,8 +7,6 @@ Invariants:
   Pydantic ``check_api_key`` validator does not re-hash it.
 """
 
-from typing import TYPE_CHECKING
-
 from litellm.constants import LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME
 from litellm.identity.context import AuditInfo, ClientInfo, IdentityContext, RequestIds
 from litellm.identity.jwt import parse_jwt_scopes
@@ -21,9 +19,7 @@ from litellm.identity.principal import (
     classify_principal_kind,
 )
 from litellm.identity.service_accounts import SERVICE_ACCOUNT_NAMES
-
-if TYPE_CHECKING:
-    from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
 
 
 def _principal_from_uak(uak: "UserAPIKeyAuth") -> Principal:
@@ -79,8 +75,6 @@ def user_api_key_auth_to_identity_context(
 def identity_context_to_user_api_key_auth(
     ctx: IdentityContext,
 ) -> "UserAPIKeyAuth":
-    from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
-
     principal = ctx.principal
     kwargs: dict = {
         "end_user_id": ctx.end_user_id,
diff --git a/litellm/identity/cache.py b/litellm/identity/cache.py
index b238d11dc56..7566ebcbd8f 100644
--- a/litellm/identity/cache.py
+++ b/litellm/identity/cache.py
@@ -17,9 +17,10 @@ from typing import TYPE_CHECKING, Optional
 from litellm.integrations.otel.model.spans import SpanRole
 from litellm.integrations.otel.runtime import traced
 
+from litellm.proxy._types import UserAPIKeyAuth
+
 if TYPE_CHECKING:
     from litellm.caching.dual_cache import DualCache
-    from litellm.proxy._types import UserAPIKeyAuth
 
 
 IDENTITY_KEY_PREFIX = "identity:v1"
@@ -86,8 +87,6 @@ class IdentityCache:
         },
     )
     async def get(self, token_hash: str) -> Optional["UserAPIKeyAuth"]:
-        from litellm.proxy._types import UserAPIKeyAuth
-
         cache_key = identity_cache_key(token_hash)
         cached = await self._cache.async_get_cache(
             key=cache_key, model_type=UserAPIKeyAuth
diff --git a/litellm/identity/extractors/api_key.py b/litellm/identity/extractors/api_key.py
index 2602e394c38..f5e311ed0d1 100644
--- a/litellm/identity/extractors/api_key.py
+++ b/litellm/identity/extractors/api_key.py
@@ -8,6 +8,7 @@ Wraps the existing ``get_api_key`` (header extraction) and
 from typing import Optional
 
 from litellm.identity.principal import ApiKeyPrincipal
+from litellm.proxy._types import UserAPIKeyAuth
 
 
 def hash_principal_token(api_key: str) -> str:
@@ -17,8 +18,6 @@ def hash_principal_token(api_key: str) -> str:
     keys) can reuse the same hashing without re-importing the Pydantic
     model.
     """
-    from litellm.proxy._types import UserAPIKeyAuth
-
     return UserAPIKeyAuth._safe_hash_litellm_api_key(api_key)
 
 
diff --git a/litellm/identity/extractors/client.py b/litellm/identity/extractors/client.py
index 87fae888a3c..a1ff1899857 100644
--- a/litellm/identity/extractors/client.py
+++ b/litellm/identity/extractors/client.py
@@ -7,6 +7,7 @@ trusted proxy (delegated to ``IPAddressUtils.is_request_from_trusted_proxy``).
 from typing import Any, Dict, List, Mapping, Optional
 
 from litellm.identity.context import ClientInfo
+from litellm.proxy.auth.ip_address_utils import IPAddressUtils
 
 
 def _split_forwarded_chain(raw: Optional[str]) -> List[str]:
@@ -34,8 +35,6 @@ def extract_client_info(
     if xff:
         forwarded_chain = _split_forwarded_chain(xff)
 
-    from litellm.proxy.auth.ip_address_utils import IPAddressUtils
-
     ip: Optional[str]
     if forwarded_chain and IPAddressUtils.is_request_from_trusted_proxy(
         request=request, general_settings=general_settings
diff --git a/litellm/identity/extractors/end_user.py b/litellm/identity/extractors/end_user.py
index db7b1c52d91..9add5852269 100644
--- a/litellm/identity/extractors/end_user.py
+++ b/litellm/identity/extractors/end_user.py
@@ -6,6 +6,8 @@ identifier only. DB validation stays in ``resolve_and_validate_end_user_id``.
 
 from typing import Optional
 
+from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body
+
 
 def extract_end_user_id(
     body: Optional[dict],
@@ -13,7 +15,4 @@ def extract_end_user_id(
 ) -> Optional[str]:
     if body is None:
         body = {}
-
-    from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body
-
     return get_end_user_id_from_request_body(request_body=body, request_headers=headers)
diff --git a/litellm/identity/jwt.py b/litellm/identity/jwt.py
index d0b9bdf07ff..96b90943b47 100644
--- a/litellm/identity/jwt.py
+++ b/litellm/identity/jwt.py
@@ -7,10 +7,9 @@ the carrier.
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Tuple
+from typing import Any, Tuple
 
-if TYPE_CHECKING:
-    from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
 
 
 def parse_jwt_scopes(claims: dict) -> Tuple[str, ...]:
@@ -45,8 +44,6 @@ def build_user_api_key_auth_from_jwt_result(
     because the auth_builder result reports it but the call site already
     branched on it; threading it through avoids re-checking.
     """
-    from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
-
     team_id = result["team_id"]
     team_object = result["team_object"]
     user_id = result["user_id"]
diff --git a/litellm/identity/oauth2.py b/litellm/identity/oauth2.py
index 016f54b2405..dc403ed33eb 100644
--- a/litellm/identity/oauth2.py
+++ b/litellm/identity/oauth2.py
@@ -6,10 +6,9 @@ an already-validated response payload into the carrier.
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Optional
+from typing import Optional
 
-if TYPE_CHECKING:
-    from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
+from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
 
 
 def build_user_api_key_auth_from_oauth2_response(
@@ -31,8 +30,6 @@ def build_user_api_key_auth_from_oauth2_response(
     Active-token validation, scope checks, and token-not-active rejection
     happen upstream in ``Oauth2Handler``; this builder trusts its input.
     """
-    from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
-
     user_id: Optional[str] = response_data.get(user_id_field_name)
     raw_role = response_data.get(user_role_field_name)
     user_team_id: Optional[str] = response_data.get(user_team_id_field_name)
diff --git a/litellm/identity/resolver.py b/litellm/identity/resolver.py
index 801d102be8a..1418ce475bb 100644
--- a/litellm/identity/resolver.py
+++ b/litellm/identity/resolver.py
@@ -23,6 +23,7 @@ from litellm.identity.principal import (
     ServiceAccountPrincipal,
 )
 from litellm.identity.service_accounts import SERVICE_ACCOUNT_NAMES
+from litellm.identity.store import load_identity
 from litellm.integrations.otel.model.spans import SpanRole
 from litellm.integrations.otel.runtime import traced
 
@@ -104,8 +105,6 @@ async def resolve_user_api_key_auth(
     centralized via ``hash_principal_token`` so the auth chain doesn't
     re-implement the JWT-vs-key hashing rules.
     """
-    from litellm.identity.store import load_identity
-
     hashed_token = hash_principal_token(api_key)
     return await load_identity(
         hashed_token=hashed_token,
diff --git a/litellm/identity/store.py b/litellm/identity/store.py
index 4d9e6e2f8a3..bb732c6872e 100644
--- a/litellm/identity/store.py
+++ b/litellm/identity/store.py
@@ -13,10 +13,20 @@ from typing import TYPE_CHECKING, Optional
 from litellm.identity.principal import classify_principal_kind
 from litellm.integrations.otel.model.spans import SpanRole
 from litellm.integrations.otel.runtime import traced
-from litellm.proxy._types import ProxyErrorTypes, ProxyException
+from litellm.proxy._types import (
+    LiteLLM_UserTable,
+    ProxyErrorTypes,
+    ProxyException,
+    UserAPIKeyAuth,
+)
+from litellm.proxy.auth.auth_checks import (
+    _cache_key_object,
+    _fetch_key_object_from_db_with_reconnect,
+    get_object_permission,
+    get_user_object,
+)
 
 if TYPE_CHECKING:
-    from litellm.proxy._types import UserAPIKeyAuth
     from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
     from litellm.proxy.utils import PrismaClient, ProxyLogging
 
@@ -39,13 +49,6 @@ async def _fetch_from_db(
     parent_otel_span,
     proxy_logging_obj: Optional["ProxyLogging"],
 ) -> Optional["UserAPIKeyAuth"]:
-    from litellm.proxy._types import UserAPIKeyAuth
-    from litellm.proxy.auth.auth_checks import (
-        _fetch_key_object_from_db_with_reconnect,
-        get_object_permission,
-        get_user_object,
-    )
-
     row = await _fetch_key_object_from_db_with_reconnect(
         hashed_token=hashed_token,
         prisma_client=prisma_client,
@@ -170,8 +173,6 @@ async def _populate_legacy_cache(
     populating it on every cold load preserves that fast path with no
     code change at the call site.
     """
-    from litellm.proxy.auth.auth_checks import _cache_key_object
-
     try:
         await _cache_key_object(
             hashed_token=hashed_token,
@@ -191,8 +192,6 @@ def _rehydrate_bundled_user(uak: "UserAPIKeyAuth") -> None:
     `metadata`, `user_role`, …) off the model, so we restore the typed
     form before handing the cached entry back.
     """
-    from litellm.proxy._types import LiteLLM_UserTable
-
     raw = getattr(uak, "user", None)
     if isinstance(raw, dict):
         try:
diff --git a/tests/test_litellm/identity/test_resolver.py b/tests/test_litellm/identity/test_resolver.py
index 389333ef5f9..fddba27886f 100644
--- a/tests/test_litellm/identity/test_resolver.py
+++ b/tests/test_litellm/identity/test_resolver.py
@@ -85,7 +85,7 @@ async def test_resolve_user_api_key_auth_hashes_token_then_loads_identity():
 
     sentinel = object()
     with patch(
-        "litellm.identity.store.load_identity",
+        "litellm.identity.resolver.load_identity",
         new=AsyncMock(return_value=sentinel),
     ) as mock_load:
         result = await resolve_user_api_key_auth(
diff --git a/tests/test_litellm/identity/test_store.py b/tests/test_litellm/identity/test_store.py
index aece8e7713a..93f2705542c 100644
--- a/tests/test_litellm/identity/test_store.py
+++ b/tests/test_litellm/identity/test_store.py
@@ -172,7 +172,7 @@ async def test_cold_path_bundles_user_into_cache():
     with (
         patch("litellm.identity.store._populate_legacy_cache", new=AsyncMock()),
         patch(
-            "litellm.proxy.auth.auth_checks.get_user_object",
+            "litellm.identity.store.get_user_object",
             side_effect=_fake_get_user_object,
         ),
     ):
@@ -254,7 +254,7 @@ async def test_cold_path_fetches_object_permission_when_only_id_present():
     with (
         patch("litellm.identity.store._populate_legacy_cache", new=AsyncMock()),
         patch(
-            "litellm.proxy.auth.auth_checks.get_object_permission",
+            "litellm.identity.store.get_object_permission",
             new=AsyncMock(return_value=fetched),
         ) as mock_get_perm,
     ):
@@ -285,7 +285,7 @@ async def test_cold_path_swallows_object_permission_lookup_failure():
     with (
         patch("litellm.identity.store._populate_legacy_cache", new=AsyncMock()),
         patch(
-            "litellm.proxy.auth.auth_checks.get_object_permission",
+            "litellm.identity.store.get_object_permission",
             new=AsyncMock(side_effect=RuntimeError("db blip")),
         ),
     ):
@@ -307,7 +307,7 @@ async def test_populate_legacy_cache_delegates_to_cache_key_object():
 
     uak = UserAPIKeyAuth(token="hash-legacy", user_id="u1")
     with patch(
-        "litellm.proxy.auth.auth_checks._cache_key_object", new=AsyncMock()
+        "litellm.identity.store._cache_key_object", new=AsyncMock()
     ) as mock_cache_key:
         await _populate_legacy_cache(
             hashed_token="hash-legacy",
@@ -328,7 +328,7 @@ async def test_populate_legacy_cache_swallows_write_failures():
 
     uak = UserAPIKeyAuth(token="hash-legacy", user_id="u1")
     with patch(
-        "litellm.proxy.auth.auth_checks._cache_key_object",
+        "litellm.identity.store._cache_key_object",
         new=AsyncMock(side_effect=RuntimeError("redis down")),
     ):
         await _populate_legacy_cache(