scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

refactor: module level imports for ensure_litellm_metadata and CodeQL

Browse source
This commit is contained in:
Alexander Grattan 2026-04-13 12:57:38 -04:00
parent c371536169
commit 2d4ea48fec
2 changed files with 20 additions and 19 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py
View file

@ -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
)

32
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_grayswan.py
View file

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