diff --git a/litellm/proxy/common_utils/responses_stream_errors.py b/litellm/proxy/common_utils/responses_stream_errors.py index 7bec25bc941..88ef6c25422 100644 --- a/litellm/proxy/common_utils/responses_stream_errors.py +++ b/litellm/proxy/common_utils/responses_stream_errors.py @@ -1,12 +1,12 @@ import json import time -from collections.abc import Mapping +from collections.abc import Iterable, Mapping, Sequence from http import HTTPStatus from types import MappingProxyType from typing import Final from pydantic import BaseModel, ConfigDict, ValidationError, field_validator -from typing_extensions import TypeIs # noqa: TID251 # TypeIs reaches typing only on 3.13 +from typing_extensions import TypeIs # noqa: TID251 # exact-type checks keep subclasses out, isinstance cannot from litellm._logging import redact_internal_details_from_client_message from litellm._uuid import uuid @@ -68,43 +68,48 @@ class _ErrorPayload(LiteLLMBaseModel): return {"message": self.error} if isinstance(self.error, str) else self.error -def _is_dict_like( - value: object, -) -> TypeIs[dict[object, object] | MappingProxyType[object, object]]: # guard-ok: isinstance decides it - return isinstance(value, dict | MappingProxyType) +def _is_plain_mapping(value: object) -> TypeIs[Mapping[object, object]]: # guard-ok: exact type check + return type(value) in (dict, MappingProxyType) -def _is_mapping_proxy( - value: object, -) -> TypeIs[MappingProxyType[object, object]]: # guard-ok: isinstance decides it, keys and values stay object - return isinstance(value, MappingProxyType) +def _is_plain_sequence(value: object) -> TypeIs[Sequence[object]]: # guard-ok: exact type check + return type(value) in (list, tuple) def _json_key(key: object) -> str | int | float | bool | None: - return key if key is None or isinstance(key, str | int | float | bool) else str(key) + return key if key is None or isinstance(key, str | int | float | bool) else repr(key) -def _with_json_keys(value: Mapping[object, object]) -> Mapping[str | int | float | bool | None, object]: - return {_json_key(key): item for key, item in value.items()} +def _json_safe_items(items: Iterable[tuple[object, object]]) -> Mapping[str | int | float | bool | None, object]: + return {_json_key(key): _json_safe(item) for key, item in items} -def _jsonable_stream_value(value: object) -> object: - if _is_mapping_proxy(value): - return _with_json_keys(value) - if isinstance(value, BaseModel): - return value.model_dump() - return str(value) +def _json_safe(value: object) -> object: + if value is None or isinstance(value, str | int | float | bool): + return value + if _is_plain_mapping(value): + return _json_safe_items(value.items()) + if _is_plain_sequence(value): + return [_json_safe(item) for item in value] + if isinstance(value, BaseModel) and type(value).__repr__ is BaseModel.__repr__: + shown: Final[Iterable[tuple[str | None, object]]] = value.__repr_args__() + return _json_safe_items(shown) + return repr(value) -def _stream_json(value: Mapping[object, object]) -> str | None: +def _stream_json(value: object) -> str | None: try: - return json.dumps(_with_json_keys(value), default=_jsonable_stream_value) - except Exception: # noqa: BLE001 # hook values run their own __str__ and serializers, so any failure must fall back + return json.dumps(_json_safe(value)) + except Exception: # noqa: BLE001 # hook values run their own __repr__, so any failure must fall back return None +def _encodable_items(error: Mapping[str, object]) -> Mapping[str, object]: + return {key: value for key, value in error.items() if _stream_json(value) is not None} + + def responses_chunk_json(chunk: object) -> str | None: - return _stream_json(chunk) if _is_dict_like(chunk) else None + return _stream_json(chunk) if _is_plain_mapping(chunk) else None def sse_error_object(data: str | bytes) -> Mapping[str, object] | None: @@ -192,6 +197,18 @@ def responses_error_fields(error: Mapping[str, object]) -> Mapping[str, object]: } +def _error_event( + fields: Mapping[str, object], error: Mapping[str, object], sequence_number: int +) -> Mapping[str, object]: + return { + "type": ResponsesAPIStreamEvents.ERROR.value, + **fields, + "param": error.get("param"), + "sequence_number": sequence_number, + "error": error, + } + + class ResponsesStreamErrorState: def __init__(self) -> None: self.response_id: str | None = None @@ -225,22 +242,8 @@ class ResponsesStreamErrorState: def format_error(self, error: Mapping[str, object]) -> str: fields: Final = responses_error_fields(error) sequence_number: Final = self.sequence_number + 1 - payload: Final = _stream_json( - { - "type": ResponsesAPIStreamEvents.ERROR.value, - **fields, - "param": error.get("param"), - "sequence_number": sequence_number, - "error": error, - } - ) or json.dumps( - { - "type": ResponsesAPIStreamEvents.ERROR.value, - **fields, - "param": None, - "sequence_number": sequence_number, - "error": fields, - } + payload: Final = _stream_json(_error_event(fields, _encodable_items(error), sequence_number)) or json.dumps( + _error_event(fields, fields, sequence_number) ) return f"event: error\ndata: {payload}\n\n" diff --git a/tests/unit/proxy/proxy_server/test_streaming_helpers.py b/tests/unit/proxy/proxy_server/test_streaming_helpers.py index ab0a2f3b6ba..4dfb3daec52 100644 --- a/tests/unit/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/unit/proxy/proxy_server/test_streaming_helpers.py @@ -27,7 +27,7 @@ import pytest from fastapi import HTTPException, Response from fastapi.responses import JSONResponse, StreamingResponse from openai import APIError as OpenAIAPIError -from pydantic import BaseModel +from pydantic import BaseModel, Field import litellm import litellm.proxy.proxy_server as ps @@ -61,6 +61,7 @@ from litellm.types.llms.openai import ( ResponseIncompleteEvent, ResponsesAPIResponse, ) +from litellm.types.proxy.litellm_pre_call_utils import RedactedDict from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage from tests._support.clock_loops import AutojumpClockLoop @@ -235,12 +236,43 @@ def _responses_frame_payload(nested: object) -> Mapping[str, object]: return chunk_payload -def test_responses_stream_frames_send_nested_header_mappings_in_their_masked_string_form() -> None: - assert _responses_frame_payload(_UPSTREAM_HEADERS)["nested"] == str(_UPSTREAM_HEADERS) +class _HiddenFieldModel(BaseModel, frozen=True): + user: str + password: str = Field(repr=False) -def test_responses_stream_frames_turn_mapping_keys_json_cannot_take_into_strings() -> None: - assert _responses_frame_payload(MappingProxyType({("a", "b"): 1}))["nested"] == {str(("a", "b")): 1} +class _MaskedReprModel(BaseModel, frozen=True): + token: str + + def __repr__(self) -> str: + return "_MaskedReprModel(token=[secure])" + + +_REDACTED_HEADERS: Final = RedactedDict({"authorization": "Bearer hook-canary"}) + + +@pytest.mark.parametrize( + ("nested", "expected"), + [ + (_UPSTREAM_HEADERS, repr(_UPSTREAM_HEADERS)), + (_REDACTED_HEADERS, repr(_REDACTED_HEADERS)), + (MappingProxyType({"headers": _REDACTED_HEADERS}), {"headers": repr(_REDACTED_HEADERS)}), + (_HiddenFieldModel(user="rig", password="hook-canary"), {"user": "rig"}), + (_MaskedReprModel(token="hook-canary"), repr(_MaskedReprModel(token="hook-canary"))), + ], + ids=["header_mapping", "redacted_dict", "proxy_over_redacted_dict", "repr_false_field", "masked_repr_model"], +) +def test_responses_stream_frames_never_show_what_a_nested_value_hides_from_its_repr( + nested: object, expected: object +) -> None: + assert _responses_frame_payload(nested)["nested"] == expected + + +@pytest.mark.parametrize("mapping", [dict, MappingProxyType], ids=["dict", "mapping_proxy"]) +def test_responses_stream_frames_turn_mapping_keys_json_cannot_take_into_strings( + mapping: Callable[[Mapping[object, object]], Mapping[object, object]], +) -> None: + assert _responses_frame_payload(mapping({("a", "b"): 1}))["nested"] == {repr(("a", "b")): 1} def test_responses_sse_chunk_turns_top_level_keys_json_cannot_take_into_strings() -> None: @@ -249,7 +281,7 @@ def test_responses_sse_chunk_turns_top_level_keys_json_cannot_take_into_strings( ) assert isinstance(frame, str) - assert json.loads(frame.removeprefix("data: ")) == {"type": "response.output_text.delta", str(("a", "b")): "x"} + assert json.loads(frame.removeprefix("data: ")) == {"type": "response.output_text.delta", repr(("a", "b")): "x"} def _self_referencing_chunk() -> Mapping[str, object]: @@ -259,19 +291,18 @@ def _self_referencing_chunk() -> Mapping[str, object]: class _UnprintableValue: - def __str__(self) -> str: + def __repr__(self) -> str: raise RuntimeError("no string form") @pytest.mark.parametrize( "chunk", [ - {"type": "response.output_text.delta", "meta": {("a", "b"): 1}}, _self_referencing_chunk(), - {"type": "response.output_text.delta", "meta": _UnprintableValue()}, _UPSTREAM_HEADERS, + RedactedDict({"type": "response.output_text.delta", "authorization": "Bearer hook-canary"}), ], - ids=["nested_dict_tuple_key", "self_reference", "unprintable_value", "header_mapping_chunk"], + ids=["self_reference", "header_mapping_chunk", "redacted_dict_chunk"], ) def test_responses_sse_chunk_json_cannot_encode_streams_as_a_plain_data_frame(chunk: Mapping[str, object]) -> None: assert _format_responses_sse_chunk(chunk=chunk) == _format_streaming_sse_chunk(chunk=chunk) @@ -279,20 +310,18 @@ def test_responses_sse_chunk_json_cannot_encode_streams_as_a_plain_data_frame(ch @pytest.mark.parametrize( "nested", - [ - MappingProxyType({"bad": {("a", "b"): 1}}), - _self_referencing_chunk(), - _UnprintableValue(), - ], - ids=["proxy_over_tuple_key", "self_reference", "unprintable_value"], + [_self_referencing_chunk(), _UnprintableValue()], + ids=["self_reference", "unprintable_value"], ) -def test_responses_error_event_json_cannot_encode_still_ends_the_stream_with_its_message(nested: object) -> None: - frame: Final = ResponsesStreamErrorState().format_error({"message": "blocked", "code": 400, "nested": nested}) +def test_responses_error_event_drops_only_the_error_detail_json_cannot_encode(nested: object) -> None: + frame: Final = ResponsesStreamErrorState().format_error( + {"message": "blocked", "code": "400", "type": "guardrail_error", "nested": nested} + ) assert frame.startswith("event: error\ndata: ") payload: Final = json.loads(frame.split("data: ", 1)[1]) - assert (payload["type"], payload["message"], payload["error"]["message"]) == ("error", "blocked", "blocked") - assert payload["error"]["code"] == payload["code"] + assert (payload["type"], payload["message"]) == ("error", "blocked") + assert payload["error"] == {"message": "blocked", "code": "400", "type": "guardrail_error"} @pytest.mark.asyncio @pytest.mark.parametrize(