fix(proxy): read hook container storage the way repr does and send error-event models as str()

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)

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(