From 5ea92486a668908006f5b9e370746dd519b0d2a9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 13:58:45 -0700 Subject: [PATCH] fix(types): keep loose payloads working under the sweep's validators An off-schema Anthropic stream event is now skipped instead of ending the guardrail's text and stop_reason scan, and MAI image usage accepts the loose numbers it did before. New dict and list adapters hide their input from validation errors, RerankResponse indexing keeps its field types, and the delete_user test mocks the int /user/delete returns. --- litellm/caching/caching.py | 4 +++- .../a2a/chat/guardrail_translation/handler.py | 2 +- .../chat/guardrail_translation/handler.py | 4 ++-- .../image_generation/mai_transformation.py | 6 ++--- .../audio_transcription/transformation.py | 2 +- .../dall_e_2_transformation.py | 2 +- .../dall_e_3_transformation.py | 2 +- .../image_generation/gpt_transformation.py | 2 +- litellm/proxy/auth/oauth2_check.py | 6 ++--- litellm/proxy/client/users.py | 4 ++-- litellm/types/rerank.py | 14 +++++++++++- .../test_anthropic_guardrail_handler.py | 22 +++++++++++++++++++ .../test_mai_image_generation.py | 8 +++++++ tests/unit/proxy/client/test_users.py | 5 ++--- 14 files changed, 63 insertions(+), 20 deletions(-) diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 7ee0dc87a88..91863683dd9 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -44,7 +44,9 @@ from .s3_cache import S3Cache RESPONSE_CACHE_TARGET: Final = "llm_response" _ALWAYS_KEYED_PROVIDER_PARAMS: Final = frozenset({"continuation_token"}) _DECODED_LITERAL: Final = TypeAdapter(object) -_REDIS_STARTUP_NODES_ADAPTER: Final = TypeAdapter(list[dict[str, object]], config=ConfigDict(strict=True)) +_REDIS_STARTUP_NODES_ADAPTER: Final = TypeAdapter( + list[dict[str, object]], config=ConfigDict(strict=True, hide_input_in_errors=True) +) @contextmanager diff --git a/litellm/llms/a2a/chat/guardrail_translation/handler.py b/litellm/llms/a2a/chat/guardrail_translation/handler.py index e68f8699c4c..921d33f16b6 100644 --- a/litellm/llms/a2a/chat/guardrail_translation/handler.py +++ b/litellm/llms/a2a/chat/guardrail_translation/handler.py @@ -29,7 +29,7 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth -_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) +_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True)) class _A2ATextPart(TypedDict, total=False): diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index cb69d4f0766..01937f343ee 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -1823,7 +1823,7 @@ class AnthropicMessagesHandler(BaseTranslation): delta: _SSETextDelta = data.get("delta", {}) if delta.get("type") == "text_delta": text += delta.get("text", "") - except json.JSONDecodeError: + except (json.JSONDecodeError, ValidationError): verbose_proxy_logger.warning("Failed to parse JSON from SSE data: %s", data_line) except Exception as e: @@ -1887,7 +1887,7 @@ class AnthropicMessagesHandler(BaseTranslation): stop_reason: object = delta.get("stop_reason") if stop_reason is not None: return True - except json.JSONDecodeError: + except (json.JSONDecodeError, ValidationError): verbose_proxy_logger.warning("Failed to parse JSON from SSE data: %s", data_line) except Exception as e: diff --git a/litellm/llms/azure_ai/image_generation/mai_transformation.py b/litellm/llms/azure_ai/image_generation/mai_transformation.py index 7bb440941bd..de6890cdbcc 100644 --- a/litellm/llms/azure_ai/image_generation/mai_transformation.py +++ b/litellm/llms/azure_ai/image_generation/mai_transformation.py @@ -21,13 +21,13 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer -@with_config(ConfigDict(extra="allow", strict=True)) +@with_config(ConfigDict(extra="allow")) class _MAIImageUsageDetails(TypedDict, total=False): text_tokens: ReadOnly[NotRequired[int | float | None]] image_tokens: ReadOnly[NotRequired[int | float | None]] -@with_config(ConfigDict(extra="allow", strict=True)) +@with_config(ConfigDict(extra="allow")) class _MAIImageUsage(TypedDict, total=False): input_tokens_details: ReadOnly[NotRequired[object]] num_input_text_tokens: ReadOnly[NotRequired[int | float | None]] @@ -39,7 +39,7 @@ class _MAIImageUsage(TypedDict, total=False): total_tokens: ReadOnly[NotRequired[int | float | None]] -@with_config(ConfigDict(extra="allow", strict=True)) +@with_config(ConfigDict(extra="allow")) class _MAIImageResponse(TypedDict, total=False): usage: ReadOnly[NotRequired[object]] diff --git a/litellm/llms/edenai/audio_transcription/transformation.py b/litellm/llms/edenai/audio_transcription/transformation.py index 8430e2c3cdc..9ff530389e2 100644 --- a/litellm/llms/edenai/audio_transcription/transformation.py +++ b/litellm/llms/edenai/audio_transcription/transformation.py @@ -22,7 +22,7 @@ from litellm.utils import convert_to_model_response_object from ..common_utils import EdenAIException, authorized_headers, endpoint_url, reported_cost -_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) +_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True)) def _form_fields(model: str, optional_params: Mapping[str, object]) -> dict[str, object]: # mutable-ok: httpx form data diff --git a/litellm/llms/openai/image_generation/dall_e_2_transformation.py b/litellm/llms/openai/image_generation/dall_e_2_transformation.py index 47908f20d22..eaacb4031b8 100644 --- a/litellm/llms/openai/image_generation/dall_e_2_transformation.py +++ b/litellm/llms/openai/image_generation/dall_e_2_transformation.py @@ -10,7 +10,7 @@ from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams from litellm.types.utils import ImageResponse from litellm.utils import convert_to_model_response_object -_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) +_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True)) if TYPE_CHECKING: from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj diff --git a/litellm/llms/openai/image_generation/dall_e_3_transformation.py b/litellm/llms/openai/image_generation/dall_e_3_transformation.py index dbab449ca25..ca6c6580c9d 100644 --- a/litellm/llms/openai/image_generation/dall_e_3_transformation.py +++ b/litellm/llms/openai/image_generation/dall_e_3_transformation.py @@ -10,7 +10,7 @@ from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams from litellm.types.utils import ImageResponse from litellm.utils import convert_to_model_response_object -_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) +_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True)) if TYPE_CHECKING: from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj diff --git a/litellm/llms/openai/image_generation/gpt_transformation.py b/litellm/llms/openai/image_generation/gpt_transformation.py index 88a0b8cc01f..c29fe6a7d4f 100644 --- a/litellm/llms/openai/image_generation/gpt_transformation.py +++ b/litellm/llms/openai/image_generation/gpt_transformation.py @@ -10,7 +10,7 @@ from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams from litellm.types.utils import ImageResponse from litellm.utils import convert_to_model_response_object -_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) +_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True)) if TYPE_CHECKING: from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj diff --git a/litellm/proxy/auth/oauth2_check.py b/litellm/proxy/auth/oauth2_check.py index ece41a7de82..641e39045c1 100644 --- a/litellm/proxy/auth/oauth2_check.py +++ b/litellm/proxy/auth/oauth2_check.py @@ -3,7 +3,7 @@ import os from typing import Final, cast import httpx -from pydantic import TypeAdapter +from pydantic import ConfigDict, TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.llms.custom_httpx.http_handler import ( @@ -12,7 +12,7 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth -_TOKEN_INFO_RESPONSE: Final = TypeAdapter(dict[str, object]) +_TOKEN_INFO_RESPONSE: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True)) class Oauth2Handler: @@ -182,7 +182,7 @@ class Oauth2Handler: response.raise_for_status() # If we get here, the request was successful - data = _TOKEN_INFO_RESPONSE.validate_python(response.json(), strict=True) + data = _TOKEN_INFO_RESPONSE.validate_python(response.json()) verbose_proxy_logger.debug( "Oauth2 token validation for token=%s, response from endpoint=%s", diff --git a/litellm/proxy/client/users.py b/litellm/proxy/client/users.py index 20b380018fb..11a32896a99 100644 --- a/litellm/proxy/client/users.py +++ b/litellm/proxy/client/users.py @@ -6,8 +6,8 @@ from pydantic import ConfigDict, TypeAdapter from .exceptions import NotFoundError, UnauthorizedError -_JSON_VALUE: Final = TypeAdapter(object, config=ConfigDict(strict=True)) -_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) +_JSON_VALUE: Final = TypeAdapter(object, config=ConfigDict(strict=True, hide_input_in_errors=True)) +_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True)) class UsersManagementClient: diff --git a/litellm/types/rerank.py b/litellm/types/rerank.py index 84e787f4ee1..3e043a6fdc6 100644 --- a/litellm/types/rerank.py +++ b/litellm/types/rerank.py @@ -4,7 +4,7 @@ https://docs.cohere.com/reference/rerank """ -from typing import Literal +from typing import Literal, overload from pydantic import ConfigDict, PrivateAttr from typing_extensions import ReadOnly, Required, TypedDict @@ -95,6 +95,18 @@ class RerankResponse(LiteLLMBaseModel): def hidden_params(self, hidden_params: dict[str, object]) -> None: # mutable-ok: API requires mutation self._hidden_params = hidden_params + @overload + def __getitem__(self, key: Literal["id"]) -> str | None: ... + + @overload + def __getitem__(self, key: Literal["results"]) -> list[RerankResponseResult] | None: ... + + @overload + def __getitem__(self, key: Literal["meta"]) -> RerankResponseMeta | None: ... + + @overload + def __getitem__(self, key: str) -> object: ... + def __getitem__(self, key: str) -> object: return self.__dict__[key] 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 9e88b59adc9..a2aaacffed6 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 @@ -302,6 +302,28 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing: texts.append(data["delta"]["text"]) return texts + def test_off_schema_delta_is_skipped_without_ending_the_text_scan(self): + events = [ + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "abc"}}, + {"type": "content_block_delta", "index": 1, "delta": {"type": "thinking_delta", "text": None}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "def"}}, + ] + chunk = "".join(f"event: content_block_delta\ndata: {json.dumps(event)}\n\n" for event in events).encode() + + assert AnthropicMessagesHandler().get_streaming_string_so_far([chunk]) == "abcdef" + + def test_off_schema_message_delta_does_not_hide_a_later_stop_reason(self): + events = [ + {"type": "message_delta", "delta": None}, + {"type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": None}}, + ] + chunk = "".join(f"event: message_delta\ndata: {json.dumps(event)}\n\n" for event in events).encode() + + scan_key = AnthropicMessagesHandler().get_streaming_scan_key([chunk]) + + assert scan_key is not None + assert scan_key.stream_ended is True + @pytest.mark.asyncio async def test_deliver_ended_stream_rewrites_writes_text_back_into_sse_chunks(self): handler = AnthropicMessagesHandler() 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 27e78d35c69..3f2bee97da2 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 @@ -386,6 +386,14 @@ class TestAzureMAIImageGeneration: assert normalized["input_tokens"] == 22 assert normalized["total_tokens"] == 22 + def test_normalize_mai_usage_accepts_a_bool_token_count(self): + normalized = AzureFoundryMAIImageGenerationConfig().normalize_mai_image_usage( + {"input_tokens": 22, "output_tokens": True} + ) + + assert normalized["output_tokens"] == 1 + assert normalized["total_tokens"] == 23 + 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/client/test_users.py b/tests/unit/proxy/client/test_users.py index 5b4d89420ab..a936de77c3e 100644 --- a/tests/unit/proxy/client/test_users.py +++ b/tests/unit/proxy/client/test_users.py @@ -5,7 +5,6 @@ import pytest import requests - from litellm.proxy.client.users import ( NotFoundError, UnauthorizedError, @@ -72,9 +71,9 @@ def test_create_user_unauthorized(mock_post, client): @patch("requests.post") def test_delete_user_success(mock_post, client): mock_post.return_value.status_code = 200 - mock_post.return_value.json.return_value = {"deleted": 1} + mock_post.return_value.json.return_value = 1 result = client.delete_user(["u1"]) - assert result["deleted"] == 1 + assert result == 1 mock_post.assert_called_once()