From 2d4ea48fecc56ac6874c76788be747b9f47bcc5c Mon Sep 17 00:00:00 2001
From: Alexander Grattan <alex@grattan.me>
Date: Mon, 13 Apr 2026 12:57:38 -0400
Subject: [PATCH] refactor: module level imports for ensure_litellm_metadata
 and CodeQL

---
 .../unified_guardrail/unified_guardrail.py    |  7 ++--
 .../guardrail_hooks/test_grayswan.py          | 32 +++++++++----------
 2 files changed, 20 insertions(+), 19 deletions(-)

diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py
index 11bc9b96fd7..367e6b2f150 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py
@@ -19,9 +19,6 @@ from litellm.integrations.custom_guardrail import CustomGuardrail
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route
 from litellm.llms import load_guardrail_translation_mappings
-from litellm.llms.base_llm.guardrail_translation.base_translation import (
-    BaseTranslation,
-)
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.types.guardrails import GuardrailEventHooks
 from litellm.types.utils import CallTypes, CallTypesLiteral
@@ -58,6 +55,10 @@ endpoint_guardrail_translation_mappings = None
 def _ensure_litellm_metadata(data: dict, user_api_key_dict: UserAPIKeyAuth) -> None:
     """Populate data['litellm_metadata'] from user_api_key_dict if absent."""
     if "litellm_metadata" not in data:
+        from litellm.llms.base_llm.guardrail_translation.base_translation import (
+            BaseTranslation,
+        )
+
         user_metadata = BaseTranslation.transform_user_api_key_dict_to_metadata(
             user_api_key_dict
         )
diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_grayswan.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_grayswan.py
index 4d5a7b4796c..14cf732aa7c 100644
--- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_grayswan.py
+++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_grayswan.py
@@ -4,9 +4,6 @@ import pytest
 from fastapi import HTTPException
 
 from litellm.integrations.custom_guardrail import ModifyResponseException
-from litellm.llms.openai.chat.guardrail_translation.handler import (
-    OpenAIChatCompletionsHandler,
-)
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.guardrails.guardrail_hooks.grayswan.grayswan import (
     GraySwanGuardrail,
@@ -480,13 +477,16 @@ def test_prepare_payload_includes_litellm_metadata(
     assert payload["litellm_metadata"]["user_api_key_team_id"] == "team-456"
 
 
-def test_ensure_litellm_metadata_populates_from_user_api_key_auth() -> None:
-    """Verify BaseTranslation._ensure_litellm_metadata extracts from metadata."""
-    handler = OpenAIChatCompletionsHandler()
-    user_auth = UserAPIKeyAuth(user_id="u1", team_id="t1", api_key="sk-test-hashed")
-    data: dict = {"metadata": {"user_api_key_auth": user_auth}}
+def test_ensure_litellm_metadata_populates_from_user_api_key_dict() -> None:
+    """Verify _ensure_litellm_metadata populates from user_api_key_dict."""
+    from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
+        _ensure_litellm_metadata,
+    )
 
-    handler._ensure_litellm_metadata(data)
+    user_auth = UserAPIKeyAuth(user_id="u1", team_id="t1", api_key="sk-test-hashed")
+    data: dict = {}
+
+    _ensure_litellm_metadata(data, user_auth)
 
     assert "litellm_metadata" in data
     assert data["litellm_metadata"]["user_api_key_user_id"] == "u1"
@@ -495,13 +495,13 @@ def test_ensure_litellm_metadata_populates_from_user_api_key_auth() -> None:
 
 def test_ensure_litellm_metadata_noop_when_already_present() -> None:
     """Verify _ensure_litellm_metadata does not overwrite existing litellm_metadata."""
-    handler = OpenAIChatCompletionsHandler()
-    user_auth = UserAPIKeyAuth(user_id="should-not-appear")
-    data: dict = {
-        "litellm_metadata": {"existing": "value"},
-        "metadata": {"user_api_key_auth": user_auth},
-    }
+    from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
+        _ensure_litellm_metadata,
+    )
 
-    handler._ensure_litellm_metadata(data)
+    user_auth = UserAPIKeyAuth(user_id="should-not-appear")
+    data: dict = {"litellm_metadata": {"existing": "value"}}
+
+    _ensure_litellm_metadata(data, user_auth)
 
     assert data["litellm_metadata"] == {"existing": "value"}