scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-11 08:38:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 280

refactor(types): tighten round 2 parse sites and keep payloads out of errors

Browse source 
Parse Anthropic SSE events with validate_json and ignore extra fields, validate admin team and key rows per row with hidden inputs, type rerank get and __getitem__ per key, and move the Phoenix prompt version response to a frozen BaseModel
This commit is contained in:
mateo-berri 2026-10-29 07:46:54 -07:00
parent 06d4552045
commit 2f8eca0a5b
8 changed files with 102 additions and 25 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
litellm/integrations/arize/arize_phoenix_client.py
View file

@ -5,18 +5,15 @@ Arize Phoenix API client for fetching prompt versions from Arize Phoenix.
import urllib.parse
from typing import Final
from pydantic import ConfigDict, TypeAdapter, with_config
from typing_extensions import NotRequired, ReadOnly, TypedDict
from pydantic import BaseModel, ConfigDict
from litellm.llms.custom_httpx.http_handler import HTTPHandler
@with_config(ConfigDict(extra="allow", strict=True))
class _PromptVersionResponse(TypedDict):
data: ReadOnly[NotRequired[dict[str, object] | None]]
class _PromptVersionResponse(BaseModel):
model_config = ConfigDict(extra="ignore", strict=True, frozen=True, hide_input_in_errors=True)
_PROMPT_VERSION_RESPONSE: Final = TypeAdapter(_PromptVersionResponse)
data: dict[str, object] | None = None
def _sanitize_id(identifier: str) -> str:
@ -82,8 +79,7 @@ class ArizePhoenixClient:
response: Final = self.http_handler.get(url, headers=self.headers)
response.raise_for_status()
data: Final = _PROMPT_VERSION_RESPONSE.validate_python(response.json())
return data.get("data")
return _PromptVersionResponse.model_validate(response.json()).data
except Exception as e:
# Check if it's an HTTP error

16
litellm/llms/anthropic/chat/guardrail_translation/handler.py
View file

@ -349,23 +349,23 @@ def _tool_call_shapes(tool_calls: Sequence[object]) -> tuple[_ToolCallShape, ...
)
@with_config(ConfigDict(extra="allow", strict=True))
@with_config(ConfigDict(extra="ignore", strict=True))
class _SSETextDelta(TypedDict, total=False):
type: ReadOnly[object]
text: ReadOnly[str]
@with_config(ConfigDict(extra="allow", strict=True))
@with_config(ConfigDict(extra="ignore", strict=True))
class _SSETextDeltaEvent(TypedDict, total=False):
delta: ReadOnly[_SSETextDelta]
@with_config(ConfigDict(extra="allow", strict=True))
@with_config(ConfigDict(extra="ignore", strict=True))
class _SSEStopDelta(TypedDict, total=False):
stop_reason: ReadOnly[object]
@with_config(ConfigDict(extra="allow", strict=True))
@with_config(ConfigDict(extra="ignore", strict=True))
class _SSEStopDeltaEvent(TypedDict, total=False):
delta: ReadOnly[_SSEStopDelta]
@ -1819,11 +1819,11 @@ class AnthropicMessagesHandler(BaseTranslation):
# Only process content_block_delta events
if event_type == "content_block_delta" and data_line:
try:
data: _SSETextDeltaEvent = _SSE_TEXT_DELTA_EVENT.validate_python(json.loads(data_line))
data: _SSETextDeltaEvent = _SSE_TEXT_DELTA_EVENT.validate_json(data_line)
delta: _SSETextDelta = data.get("delta", {})
if delta.get("type") == "text_delta":
text += delta.get("text", "")
except (json.JSONDecodeError, ValidationError):
except ValidationError:
verbose_proxy_logger.warning("Failed to parse JSON from SSE data: %s", data_line)
except Exception as e:
@ -1882,12 +1882,12 @@ class AnthropicMessagesHandler(BaseTranslation):
# Check for message_delta event with stop_reason
if event_type == "message_delta" and data_line:
try:
data: _SSEStopDeltaEvent = _SSE_STOP_DELTA_EVENT.validate_python(json.loads(data_line))
data: _SSEStopDeltaEvent = _SSE_STOP_DELTA_EVENT.validate_json(data_line)
delta: _SSEStopDelta = data.get("delta", {})
stop_reason: object = delta.get("stop_reason")
if stop_reason is not None:
return True
except (json.JSONDecodeError, ValidationError):
except ValidationError:
verbose_proxy_logger.warning("Failed to parse JSON from SSE data: %s", data_line)
except Exception as e:

17
litellm/proxy/management_endpoints/internal_user_endpoints.py
View file

@ -125,11 +125,16 @@ _USER_CACHE_INVALIDATION_BATCH_SIZE: Final = 50
@with_config(ConfigDict(extra="allow", strict=True))
class _AdminTeamAndKeyRows(TypedDict):
keys: ReadOnly[Sequence[Mapping[str, object]] | None]
teams: ReadOnly[Sequence[Mapping[str, object]] | None]
keys: ReadOnly[Sequence[object] | None]
teams: ReadOnly[Sequence[object] | None]
_ADMIN_TEAM_AND_KEY_ROWS_ADAPTER: Final = TypeAdapter(list[_AdminTeamAndKeyRows])
_ADMIN_TEAM_AND_KEY_ROWS_ADAPTER: Final = TypeAdapter(
list[_AdminTeamAndKeyRows], config=ConfigDict(hide_input_in_errors=True)
)
_ADMIN_KEY_ROW_ADAPTER: Final = TypeAdapter(
dict[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True)
)
def _user_table(
@ -1209,17 +1214,17 @@ async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth):
verbose_proxy_logger.debug("results_keys: %s", results)
_keys_in_db: Final[Sequence[Mapping[str, object]]] = results[0]["keys"] or []
_keys_in_db: Final[Sequence[object]] = results[0]["keys"] or []
# cast all keys to LiteLLM_VerificationToken
keys_in_db: Final = []
for key in _keys_in_db:
key_payload = dict[str, object](key)
key_payload = _ADMIN_KEY_ROW_ADAPTER.validate_python(key)
if key_payload.get("models") is None:
key_payload["models"] = []
keys_in_db.append(LiteLLM_VerificationToken.model_validate(key_payload))
# cast all teams to LiteLLM_TeamTable
_teams_rows: Final[Sequence[Mapping[str, object]]] = results[0]["teams"] or []
_teams_rows: Final[Sequence[object]] = results[0]["teams"] or []
_teams_in_db: Final = sorted(
(LiteLLM_TeamTable.model_validate(team) for team in _teams_rows),
key=lambda x: getattr(x, "team_alias", "") or "",

17
litellm/types/rerank.py
View file

@ -4,6 +4,7 @@ https://docs.cohere.com/reference/rerank
"""
from collections.abc import Sequence
from typing import Literal, overload
from pydantic import ConfigDict, PrivateAttr
@ -105,11 +106,23 @@ class RerankResponse(LiteLLMBaseModel):
def __getitem__(self, key: Literal["meta"]) -> RerankResponseMeta | None: ...
@overload
def __getitem__(self, key: str) -> object: ...
def __getitem__(self, key: str) -> str | Sequence[RerankResponseResult] | RerankResponseMeta | None: ...
def __getitem__(self, key: str) -> object:
def __getitem__(self, key: str) -> str | Sequence[RerankResponseResult] | RerankResponseMeta | None:
return self.__dict__[key]
@overload
def get(self, key: Literal["id"], default: None = None) -> str | None: ...
@overload
def get(self, key: Literal["results"], default: None = None) -> list[RerankResponseResult] | None: ...
@overload
def get(self, key: Literal["meta"], default: None = None) -> RerankResponseMeta | None: ...
@overload
def get(self, key: str, default: object = None) -> object: ...
def get(self, key: str, default: object = None) -> object:
return self.__dict__.get(key, default)

17
tests/unit/integrations/arize/test_arize_phoenix.py
View file

@ -971,6 +971,23 @@ def test_arize_phoenix_client_get_prompt_version_rejects_traversal():
client.get_prompt_version("../../projects")
def test_arize_phoenix_client_error_does_not_echo_prompt_content():
import httpx
import respx
from litellm.integrations.arize.arize_phoenix_client import ArizePhoenixClient
client = ArizePhoenixClient(api_key="test-key", api_base="https://phoenix.test")
with respx.mock:
respx.get("https://phoenix.test/v1/prompt_versions/pv-1").mock(
return_value=httpx.Response(200, json={"data": "secret system prompt text"})
)
with pytest.raises(Exception, match="Error fetching prompt version 'pv-1'") as exc_info:
client.get_prompt_version("pv-1")
assert "secret system prompt text" not in str(exc_info.value)
@pytest.fixture()
def _vcr_outcome_gate(request, vcr):
install_live_call_probe(request, vcr)

10
tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py
View file

@ -312,6 +312,16 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing:
assert AnthropicMessagesHandler().get_streaming_string_so_far([chunk]) == "abcdef"
def test_malformed_delta_data_is_skipped_without_ending_the_text_scan(self):
text_delta = {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "abc"}}
chunk = (
f"event: content_block_delta\ndata: {json.dumps(text_delta)}\n\n"
"event: content_block_delta\ndata: {not json\n\n"
f"event: content_block_delta\ndata: {json.dumps(text_delta)}\n\n"
).encode()
assert AnthropicMessagesHandler().get_streaming_string_so_far([chunk]) == "abcabc"
def test_off_schema_message_delta_does_not_hide_a_later_stop_reason(self):
events = [
{"type": "message_delta", "delta": None},

8
tests/unit/llms/azure_ai/image_generation/test_mai_image_generation.py
View file

@ -394,6 +394,14 @@ class TestAzureMAIImageGeneration:
assert normalized["output_tokens"] == 1
assert normalized["total_tokens"] == 23
def test_normalize_mai_usage_reads_string_token_counts_as_numbers(self):
normalized = AzureFoundryMAIImageGenerationConfig().normalize_mai_image_usage(
{"input_tokens": "8", "output_tokens": "1024"}
)
assert normalized["input_tokens"] == 8
assert normalized["total_tokens"] == 1032
def test_azure_sync_image_generation_uses_mai_response_transform(self):
raw_response = MagicMock(spec=httpx.Response)
raw_response.json.return_value = {

28
tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py
View file

@ -4292,6 +4292,34 @@ async def test_get_user_info_for_proxy_admin_validates_keys_and_teams():
assert returned_key["models"] == []
@pytest.mark.asyncio
@pytest.mark.parametrize(
"raw_rows",
[
[{"teams": None, "keys": "sk-row-content-must-not-leak"}],
[{"teams": None, "keys": ["sk-row-content-must-not-leak"]}],
],
)
async def test_get_user_info_for_proxy_admin_errors_do_not_echo_row_content(raw_rows):
from unittest.mock import AsyncMock, MagicMock, patch
from pydantic import ValidationError
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.management_endpoints.internal_user_endpoints import (
_get_user_info_for_proxy_admin,
)
mock_prisma_client = MagicMock()
mock_prisma_client.db.query_raw = AsyncMock(return_value=raw_rows)
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client):
with pytest.raises(ValidationError) as exc_info:
await _get_user_info_for_proxy_admin(user_api_key_dict=UserAPIKeyAuth(user_id=None))
assert "sk-row-content-must-not-leak" not in str(exc_info.value)
def _object_permission_mocks(mocker, existing_object_permission_id=None):
"""Prisma double whose user row optionally already links a permission row."""
mock_prisma_client = mocker.MagicMock()