diff --git a/litellm/proxy/common_utils/responses_stream_errors.py b/litellm/proxy/common_utils/responses_stream_errors.py index ac4854f4014..137b1221d7d 100644 --- a/litellm/proxy/common_utils/responses_stream_errors.py +++ b/litellm/proxy/common_utils/responses_stream_errors.py @@ -74,14 +74,37 @@ def _is_mapping( return isinstance(value, Mapping) -def jsonable_stream_value(value: object) -> object: - if _is_mapping(value): - return {**value} +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 _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) + + +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 _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 responses_chunk_json(chunk: object) -> str | None: + if not _is_mapping(chunk): + return None + try: + return json.dumps(_with_json_keys(chunk), default=_jsonable_stream_value) + except (TypeError, ValueError): + return None + + def sse_error_object(data: str | bytes) -> Mapping[str, object] | None: try: return _ErrorPayload.model_validate_json(data).error_object() @@ -206,7 +229,7 @@ class ResponsesStreamErrorState: "sequence_number": self.sequence_number + 1, "error": error, }, - default=jsonable_stream_value, + default=_jsonable_stream_value, ) return f"event: error\ndata: {payload}\n\n" diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 2ddb7a9241b..7acdd582658 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -539,7 +539,7 @@ from litellm.proxy.common_utils.periodic_reload_schedule import ( ) from litellm.proxy.common_utils.proxy_state import ProxyState from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob -from litellm.proxy.common_utils.responses_stream_errors import ResponsesStreamErrorState, jsonable_stream_value +from litellm.proxy.common_utils.responses_stream_errors import ResponsesStreamErrorState, responses_chunk_json from litellm.proxy.common_utils.scheduled_job_stagger import ( apply_scheduled_job_stagger, attach_job_timing_logger, @@ -10203,9 +10203,8 @@ def _format_streaming_sse_chunk(chunk: object) -> str | bytes: def _format_responses_sse_chunk(chunk: object) -> str | bytes: - if isinstance(chunk, Mapping): - return f"data: {json.dumps(chunk, default=jsonable_stream_value)}\n\n" - return _format_streaming_sse_chunk(chunk=chunk) + data: Final = responses_chunk_json(chunk) + return _format_streaming_sse_chunk(chunk=chunk) if data is None else f"data: {data}\n\n" _SSE_FRAME_DELIMITERS: Final = ("\r\n\r\n", "\n\n", "\r\r") diff --git a/tests/unit/proxy/proxy_server/test_streaming_helpers.py b/tests/unit/proxy/proxy_server/test_streaming_helpers.py index a188326c685..05bde512608 100644 --- a/tests/unit/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/unit/proxy/proxy_server/test_streaming_helpers.py @@ -219,6 +219,53 @@ def test_responses_error_event_encodes_nested_error_fields_as_json_objects() -> payload: Final = json.loads(frame.split("data: ", 1)[1]) assert payload["error"]["provider_specific_fields"] == {"upstream": {"type": "overloaded_error"}} + +_UPSTREAM_HEADERS: Final = httpx.Headers({"authorization": "Bearer hook-canary", "x-request-id": "req_1"}) + + +def _responses_frame_payload(nested: object) -> Mapping[str, object]: + chunk_frame: Final = _format_responses_sse_chunk( + chunk=MappingProxyType({"type": "response.output_text.delta", "nested": nested}) + ) + error_frame: Final = ResponsesStreamErrorState().format_error({"message": "blocked", "nested": nested}) + assert isinstance(chunk_frame, str) + chunk_payload: Final = json.loads(chunk_frame.removeprefix("data: ")) + error_payload: Final = json.loads(error_frame.split("data: ", 1)[1]) + assert chunk_payload["nested"] == error_payload["error"]["nested"] + 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) + + +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} + + +def test_responses_sse_chunk_turns_top_level_keys_json_cannot_take_into_strings() -> None: + frame: Final = _format_responses_sse_chunk( + chunk=MappingProxyType({"type": "response.output_text.delta", ("a", "b"): "x"}) + ) + + assert isinstance(frame, str) + assert json.loads(frame.removeprefix("data: ")) == {"type": "response.output_text.delta", str(("a", "b")): "x"} + + +def _self_referencing_chunk() -> Mapping[str, object]: + chunk: Final[dict[str, object]] = {"type": "response.output_text.delta"} + chunk["self"] = chunk # mutable-ok: a self-reference needs the dict to exist before it can hold itself + return chunk + + +@pytest.mark.parametrize( + "chunk", + [{"type": "response.output_text.delta", "meta": {("a", "b"): 1}}, _self_referencing_chunk()], + ids=["nested_dict_tuple_key", "self_reference"], +) +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) + @pytest.mark.asyncio @pytest.mark.parametrize( ("guardrail_error", "error_object"),