From 2f8eca0a5b61a4a47b83307157ca5c830b16f94b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 07:46:54 -0700 Subject: [PATCH] refactor(types): tighten round 2 parse sites and keep payloads out of errors 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 --- .../arize/arize_phoenix_client.py | 14 ++++------ .../chat/guardrail_translation/handler.py | 16 +++++------ .../internal_user_endpoints.py | 17 +++++++---- litellm/types/rerank.py | 17 +++++++++-- .../integrations/arize/test_arize_phoenix.py | 17 +++++++++++ .../test_anthropic_guardrail_handler.py | 10 +++++++ .../test_mai_image_generation.py | 8 ++++++ .../test_internal_user_endpoints.py | 28 +++++++++++++++++++ 8 files changed, 102 insertions(+), 25 deletions(-) diff --git a/litellm/integrations/arize/arize_phoenix_client.py b/litellm/integrations/arize/arize_phoenix_client.py index 3c1acfd36c3..4d50389faf1 100644 --- a/litellm/integrations/arize/arize_phoenix_client.py +++ b/litellm/integrations/arize/arize_phoenix_client.py @@ -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 diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 01937f343ee..08392f4fde4 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -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: diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index cc20d51aae1..e5f6e38d9a6 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -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 "", diff --git a/litellm/types/rerank.py b/litellm/types/rerank.py index 3e043a6fdc6..ffdcf53fd7e 100644 --- a/litellm/types/rerank.py +++ b/litellm/types/rerank.py @@ -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) diff --git a/tests/unit/integrations/arize/test_arize_phoenix.py b/tests/unit/integrations/arize/test_arize_phoenix.py index eb9e1326ced..5edf6e6d379 100644 --- a/tests/unit/integrations/arize/test_arize_phoenix.py +++ b/tests/unit/integrations/arize/test_arize_phoenix.py @@ -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) diff --git a/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index a2aaacffed6..26ccef5693f 100644 --- a/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -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}, diff --git a/tests/unit/llms/azure_ai/image_generation/test_mai_image_generation.py b/tests/unit/llms/azure_ai/image_generation/test_mai_image_generation.py index 3f2bee97da2..310d7797ccb 100644 --- a/tests/unit/llms/azure_ai/image_generation/test_mai_image_generation.py +++ b/tests/unit/llms/azure_ai/image_generation/test_mai_image_generation.py @@ -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 = { diff --git a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py index cb8a73d581b..ae283ac541b 100644 --- a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py @@ -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()