diff --git a/litellm/proxy/common_utils/responses_stream_errors.py b/litellm/proxy/common_utils/responses_stream_errors.py index f4acb2858fb..951650543a7 100644 --- a/litellm/proxy/common_utils/responses_stream_errors.py +++ b/litellm/proxy/common_utils/responses_stream_errors.py @@ -1,7 +1,7 @@ import json import time from collections import Counter, OrderedDict, defaultdict -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Iterable, Iterator, Mapping, Sequence from enum import Enum from http import HTTPStatus from types import MappingProxyType @@ -78,27 +78,35 @@ class _BaseForm(Enum): DUMPED = "dumped" -def _prints_every_item(kind: type[object]) -> bool: +def _prints_every_item(kind: type[object], form: _BaseForm) -> bool: repr_owner: Final = next((owner for owner in kind.__mro__ if "__repr__" in owner.__dict__), object) - return kind.__str__ is object.__str__ and repr_owner in (dict, list, tuple, OrderedDict, defaultdict, Counter) + if form is _BaseForm.DUMPED: + return kind.__str__ is object.__str__ and repr_owner in (dict, list, tuple, OrderedDict, defaultdict, Counter) + return ( + kind.__str__ is object.__str__ + and kind.__format__ is object.__format__ + and (kind in (OrderedDict, Counter) or repr_owner in (dict, list, tuple, defaultdict)) + ) -def _is_shown_mapping(value: object) -> TypeIs[Mapping[object, object]]: # guard-ok: narrows dicts to object items +def _is_shown_dict(value: object, form: _BaseForm) -> TypeIs[Mapping[object, object]]: # guard-ok: dict items as object kind: Final = type(value) - return isinstance(value, dict) and _prints_every_item(kind) + return isinstance(value, dict) and _prints_every_item(kind, form) -def _is_shown_sequence(value: object) -> TypeIs[Sequence[object]]: # guard-ok: narrows lists and tuples to object items +def _is_shown_sequence(value: object, form: _BaseForm) -> TypeIs[Sequence[object]]: # guard-ok: items as object kind: Final = type(value) - return isinstance(value, (list, tuple)) and _prints_every_item(kind) + return isinstance(value, (list, tuple)) and _prints_every_item(kind, form) def _is_proxy(value: object) -> TypeIs[Mapping[object, object]]: # guard-ok: mappingproxy cannot be subclassed return isinstance(value, MappingProxyType) -def _prints_as_dict(proxy: Mapping[object, object]) -> bool: - return repr(proxy) == f"mappingproxy({dict(proxy)!r})" +def _printed_proxy_items(proxy: Mapping[object, object], printed: str) -> Iterable[tuple[object, object]] | None: + snapshot: Final = dict(proxy) + shown: Final = repr(snapshot) + return snapshot.items() if printed in (shown, f"mappingproxy({shown})") else None def _repr_prints_its_fields(model: BaseModel) -> bool: @@ -106,11 +114,6 @@ def _repr_prints_its_fields(model: BaseModel) -> bool: return kind.__repr__ is BaseModel.__repr__ and kind.__repr_str__ is BaseModel.__repr_str__ -def _str_skips_its_repr(model: BaseModel) -> bool: - kind: Final = type(model) - return kind.__str__ is BaseModel.__str__ and kind.__repr__ is not BaseModel.__repr__ - - def _plain_scalar(value: str | float) -> str | float: if isinstance(value, str): return str.__str__(value) @@ -130,11 +133,7 @@ def _text(value: object, form: _BaseForm) -> object: try: if isinstance(value, (str, int, float)) and (form is _BaseForm.DUMPED or _repr_shows_its_value(value)): return _plain_scalar(value) - if form is _BaseForm.PRINTED: - return repr(value) - if isinstance(value, BaseModel) and _str_skips_its_repr(value): - return repr(value) - return str(value) + return repr(value) if form is _BaseForm.PRINTED else str(value) except Exception: # noqa: BLE001 # hook values run their own __repr__ and __str__, so a failure drops the value return _DROPPED @@ -143,11 +142,31 @@ def _json_key(key: object, form: _BaseForm) -> object: return key if key is None or type(key) in (str, int, float, bool) else _text(key, form) +def _shown_sequence(value: object, form: _BaseForm) -> Iterator[object] | None: + if not _is_shown_sequence(value, form): + return None + if form is _BaseForm.PRINTED and isinstance(value, list): + return list[object].__iter__(value) + if form is _BaseForm.PRINTED and isinstance(value, tuple): + return tuple[object, ...].__iter__(value) + return iter(value) + + +def _shown_items(value: object, form: _BaseForm) -> Iterable[tuple[object, object]] | None: + if not _is_shown_dict(value, form): + return None + if form is _BaseForm.PRINTED and isinstance(value, dict): + return dict[object, object].items(value) + return value.items() + + def _expansion(value: object, form: _BaseForm) -> tuple[Iterable[tuple[object, object]], _BaseForm] | None: - if _is_shown_mapping(value): - return value.items(), form - if _is_proxy(value) and _prints_as_dict(value): - return value.items(), _BaseForm.PRINTED + shown_items: Final = _shown_items(value, form) + if shown_items is not None: + return shown_items, form + if _is_proxy(value): + proxy_items: Final = _printed_proxy_items(value, repr(value) if form is _BaseForm.PRINTED else str(value)) + return None if proxy_items is None else (proxy_items, _BaseForm.PRINTED) if form is _BaseForm.PRINTED and isinstance(value, BaseModel) and _repr_prints_its_fields(value): return value.__repr_args__(), _BaseForm.PRINTED return None @@ -160,8 +179,9 @@ def _json_safe(value: object, form: _BaseForm, path: frozenset[int]) -> object: return _text(value, form) inner_path: Final = path | frozenset((id(value),)) try: - if _is_shown_sequence(value): - shown_items: Final = [_json_safe(item, form, inner_path) for item in value] + sequence: Final = _shown_sequence(value, form) + if sequence is not None: + shown_items: Final = [_json_safe(item, form, inner_path) for item in sequence] return [item for item in shown_items if item is not _DROPPED] expansion: Final = _expansion(value, form) if expansion is None: @@ -174,7 +194,9 @@ def _json_safe(value: object, form: _BaseForm, path: frozenset[int]) -> object: def responses_chunk_json(chunk: object) -> str | None: - if not _is_shown_mapping(chunk) and not _is_proxy(chunk): + if _is_proxy(chunk) and _printed_proxy_items(chunk, format(chunk, "")) is None: + return None + if not _is_shown_dict(chunk, _BaseForm.PRINTED) and not _is_proxy(chunk): return None shown: Final = _json_safe(chunk, _BaseForm.PRINTED, frozenset()) return None if isinstance(shown, str) or shown is _DROPPED else json.dumps(shown) diff --git a/tests/unit/proxy/proxy_server/test_streaming_helpers.py b/tests/unit/proxy/proxy_server/test_streaming_helpers.py index 1b900d78336..2e3bd57d43d 100644 --- a/tests/unit/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/unit/proxy/proxy_server/test_streaming_helpers.py @@ -18,8 +18,8 @@ from __future__ import annotations import asyncio import functools import json -from collections import OrderedDict -from collections.abc import AsyncIterator, Callable, Iterable, Mapping +from collections import Counter, OrderedDict +from collections.abc import AsyncIterator, Callable, ItemsView, Iterable, Iterator, Mapping from datetime import datetime, timezone from enum import Enum from http import HTTPStatus @@ -33,7 +33,7 @@ from fastapi import HTTPException, Response from fastapi.responses import JSONResponse, StreamingResponse from openai import APIError as OpenAIAPIError from openai.types.responses import ResponseOutputMessage, ResponseOutputText -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, PrivateAttr import litellm import litellm.proxy.proxy_server as ps @@ -289,7 +289,6 @@ _MASKED_IN_EVERY_FRAME: Final = ( MappingProxyType({"headers": _REDACTED_HEADERS}), MappingProxyType(_REDACTED_HEADERS), _HiddenFieldModel(user="rig", password="hook-canary"), - _MaskedReprModel(token="hook-canary"), _MaskedReprStrModel(token="hook-canary"), _MaskedReprArgsModel(token="hook-canary"), ) @@ -299,16 +298,44 @@ _MASKED_IN_EVERY_FRAME_IDS: Final = ( "proxy_over_dict_holding_redacted_dict", "proxy_wrapping_redacted_dict", "repr_false_field", - "masked_repr_model", "masked_repr_str_model", "masked_repr_args_model", ) +class _UnmaskingDict(dict[str, str]): + def items(self) -> ItemsView[str, str]: # pyright: ignore[reportIncompatibleMethodOverride] # rig hides storage + return {key: "hook-canary" for key in dict.keys(self)}.items() + + +class _UnmaskingList(list[str]): + def __iter__(self) -> Iterator[str]: + return iter(("hook-canary",)) + + +class _MostCommonMaskedCounter(Counter[str]): + def most_common(self, n: int | None = None) -> list[tuple[str, int]]: + return [(key, count) for key, count in super().most_common(n) if key != "hook-canary"] + + @pytest.mark.parametrize( "nested", - [*_MASKED_IN_EVERY_FRAME, _MaskedReprString("hook-canary")], - ids=[*_MASKED_IN_EVERY_FRAME_IDS, "masked_repr_string"], + [ + *_MASKED_IN_EVERY_FRAME, + _MaskedReprModel(token="hook-canary"), + _MaskedReprString("hook-canary"), + _UnmaskingDict(token="enc:MASKED"), + _UnmaskingList(("enc:MASKED",)), + _MostCommonMaskedCounter(("shown", "hook-canary")), + ], + ids=[ + *_MASKED_IN_EVERY_FRAME_IDS, + "masked_repr_model", + "masked_repr_string", + "dict_items_differ_from_storage", + "list_iter_differs_from_storage", + "counter_most_common_masks", + ], ) def test_responses_sse_chunk_never_shows_what_a_nested_value_hides_from_its_repr(nested: object) -> None: frame: Final = _format_responses_sse_chunk( @@ -335,18 +362,52 @@ def test_responses_error_event_never_shows_what_a_nested_value_hides_from_its_st assert "hook-canary" not in frame +class _OwnFormatDict(dict[str, str]): + def __format__(self, format_spec: str) -> str: + return json.dumps({"type": "response.output_text.delta", "token": "MASKED"}) + + +class _MaskedStrDict(dict[str, str]): + def __str__(self) -> str: + return "{'type': 'response.output_text.delta', 'token': 'MASKED'}" + + +@pytest.mark.parametrize( + "chunk", + [ + _OwnFormatDict(type="response.output_text.delta", token="hook-canary"), + MappingProxyType(_MaskedStrDict(type="response.output_text.delta", token="hook-canary")), + ], + ids=["own_format", "proxy_over_masked_str"], +) +def test_responses_sse_chunk_sends_a_whole_chunk_as_the_text_its_printer_gave(chunk: object) -> None: + frame: Final = _format_responses_sse_chunk(chunk=chunk) + + assert isinstance(frame, str) + assert "hook-canary" not in frame + + class _RigTuple(tuple[object, ...]): pass +class _RevealingReprModel(BaseModel, frozen=True): + a: int + _note: str = PrivateAttr(default="hook-canary") + + def __repr__(self) -> str: + return f"_RevealingReprModel(note={self._note})" + + @pytest.mark.parametrize( "detail", [ OrderedDict(member=_MaskedStrModel(token="hook-canary")), _RigTuple((_MaskedStrModel(token="hook-canary"),)), _HookOutputItem(id="msg_nested", type="message"), + _RevealingReprModel(a=1), ], - ids=["dict_subclass", "tuple_subclass", "model"], + ids=["dict_subclass", "tuple_subclass", "model", "model_with_own_repr"], ) def test_responses_error_event_sends_a_detail_the_way_base_json_dumps_did(detail: object) -> None: frame: Final = ResponsesStreamErrorState().format_error(