From 6bd3699d4338193b23f8e937061cbab7c4af2e49 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 16:39:14 -0700 Subject: [PATCH] fix(responses): keep guardrailed input items and bridge stream usage intact - _write_back_structured_messages now patches only the rewritten rows back into the original input items, so reasoning items (encrypted_content), function_call ids, and web_search_call items survive a guardrail rewrite verbatim; rewrites that cannot be row-mapped fall back to the previous full conversion - the responses bridge stream snapshot restores usage hidden in _hidden_params when stream_options is unset, so converted fake streams report real input_tokens instead of 0 --- litellm/llms/custom_httpx/llm_http_handler.py | 4 +- .../guardrail_translation/handler.py | 135 +++++++++++++++++- .../streaming_iterator.py | 6 + ...test_openai_responses_guardrail_handler.py | 129 ++++++++++++++++- .../guardrail_hooks/test_headroom.py | 2 +- .../test_streaming_iterator_transformation.py | 15 ++ 6 files changed, 279 insertions(+), 12 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 2ec5377e217..fcb988705eb 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2744,7 +2744,7 @@ class BaseLLMHTTPHandler: ) if self._has_agentic_completion_hook(logging_obj): - agentic_kwargs: Final = dict(litellm_params) + agentic_kwargs: Final = dict(litellm_params) # mutable-ok: agentic hooks mutate kwargs in place final_response: Final = run_async_function( self._call_agentic_completion_hooks, response=initial_response, @@ -2931,7 +2931,7 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, ) - agentic_kwargs: Final = dict(litellm_params) + agentic_kwargs: Final = dict(litellm_params) # mutable-ok: agentic hooks mutate kwargs in place final_response: Final = await self._call_agentic_completion_hooks( response=initial_response, model=model, diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 4320edb8414..8c014feeded 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -28,7 +28,8 @@ Output: response.output is List[GenericResponseOutputItem] where each has: - text: str """ -from collections.abc import Sequence +from collections.abc import Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Union, cast from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall @@ -50,6 +51,7 @@ from litellm.types.llms.openai import ( ChatCompletionToolCallChunk, ChatCompletionToolParam, OpenAIMcpServerTool, + ResponsesAPIOptionalRequestParams, ResponsesAPIStreamEvents, ) from litellm.types.responses.main import ( @@ -81,6 +83,119 @@ class ResponsesStreamChunk(TypedDict, total=False): text: ReadOnly[str] +_PATCHABLE_ITEM_FIELDS: Final[Mapping[str, str]] = MappingProxyType( + {"function_call_output": "output", "message": "content"} +) + +_EMPTY_RESPONSES_REQUEST: Final[ResponsesAPIOptionalRequestParams] = {} + + +def _item_rewrite_field(item: Mapping[str, object]) -> str | None: + item_type: Final = item.get("type") + if item_type is None: + return "content" if "content" in item else None + if not isinstance(item_type, str): + return None + return _PATCHABLE_ITEM_FIELDS.get(item_type) + + +def _rewritten_input_item(item: Mapping[str, object], rewritten: object) -> Mapping[str, object] | None: + field: Final = _item_rewrite_field(item) + if field is None or not isinstance(rewritten, Mapping): + return None + rewritten_content: Final = rewritten.get("content") + if isinstance(item.get(field), str) and isinstance(rewritten_content, str): + return {**item, field: rewritten_content} # mutable-ok: request input items must stay JSON-plain dicts + rewritten_row: Final = cast("AllMessageValues", rewritten) # cast-ok: guardrails hand back chat-shaped rows + converted_items, _ = LiteLLMResponsesTransformationHandler().convert_chat_completion_messages_to_responses_api( + [rewritten_row] # mutable-ok: converter signature takes a list + ) + if len(converted_items) != 1 or not isinstance(converted_items[0], Mapping): + return None + first_converted: Final = cast("Mapping[str, object]", converted_items[0]) # cast-ok: isinstance-checked above + converted_value: Final = first_converted.get(field) + if converted_value is None: + return None + return {**item, field: converted_value} # mutable-ok: request input items must stay JSON-plain dicts + + +def _input_item_provenance( + raw_input: Sequence[object], + expected_messages: Sequence[object], +) -> tuple[Mapping[int, int], frozenset[int]] | None: + if not all(isinstance(item, Mapping) for item in raw_input): + return None + prefixes: Final = tuple( + LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=cast("ResponseInputParam", raw_input[:count]), # cast-ok: items checked as Mappings above + responses_api_request=_EMPTY_RESPONSES_REQUEST, + ) + for count in range(len(raw_input) + 1) + ) + if tuple(prefixes[-1]) != tuple(expected_messages): + return None + item_for_message: Final = MappingProxyType( + { + message_index: item_index + for item_index in range(len(raw_input)) + for message_index in range(len(prefixes[item_index]), len(prefixes[item_index + 1])) + } + ) + tainted: Final = frozenset( + message_index + for item_index in range(len(raw_input)) + for message_index in range(len(prefixes[item_index])) + if prefixes[item_index + 1][message_index] != prefixes[item_index][message_index] + ) + return item_for_message, tainted + + +def _patch_rewritten_rows_into_input( + data: dict, + original_messages: Sequence[object], + structured_messages: Sequence[object], +) -> bool: + raw_input: Final = data.get("input") + if not isinstance(raw_input, list) or len(original_messages) != len(structured_messages): + return False + offset: Final = 1 if data.get("instructions") else 0 + provenance: Final = _input_item_provenance(raw_input, tuple(original_messages)[offset:]) + if provenance is None: + return False + item_for_message, tainted = provenance + changed: Final = tuple( + (index, rewritten) + for index, (original, rewritten) in enumerate(zip(original_messages, structured_messages)) + if original != rewritten + ) + instruction_rewrites: Final = tuple(rewritten for index, rewritten in changed if index < offset) + rewritten_instructions: Final = ( + instruction_rewrites[0].get("content") + if instruction_rewrites and isinstance(instruction_rewrites[0], Mapping) + else None + ) + if instruction_rewrites and not isinstance(rewritten_instructions, str): + return False + body_changes: Final = tuple((index - offset, rewritten) for index, rewritten in changed if index >= offset) + if any(message_index in tainted or message_index not in item_for_message for message_index, _ in body_changes): + return False + replacements: Final = MappingProxyType( + { + item_for_message[message_index]: _rewritten_input_item( + cast("Mapping[str, object]", raw_input[item_for_message[message_index]]), # cast-ok: checked Mappings + rewritten, + ) + for message_index, rewritten in body_changes + } + ) + if len(replacements) != len(body_changes) or any(item is None for item in replacements.values()): + return False + data["input"] = [replacements.get(index, item) for index, item in enumerate(raw_input)] # mutable-ok: JSON body + if isinstance(rewritten_instructions, str): + data["instructions"] = rewritten_instructions + return True + + class OpenAIResponsesHandler(BaseTranslation): """ Handler for processing OpenAI Responses API with guardrails. @@ -155,9 +270,9 @@ class OpenAIResponsesHandler(BaseTranslation): guardrailed_structured_messages is not None and guardrailed_structured_messages is not structured_messages ): - self._write_back_structured_messages(data, guardrailed_structured_messages) + self._write_back_structured_messages(data, structured_messages or (), guardrailed_structured_messages) else: - guardrailed_texts = guardrailed_inputs.get("texts", []) + guardrailed_texts = guardrailed_inputs.get("texts") or () data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data self._apply_guardrailed_tools_to_data(data, original_tools, guardrailed_inputs.get("tools")) verbose_proxy_logger.debug("OpenAI Responses API: Processed string input") @@ -217,12 +332,12 @@ class OpenAIResponsesHandler(BaseTranslation): guardrailed_structured_messages is not None and guardrailed_structured_messages is not structured_messages ): - self._write_back_structured_messages(data, guardrailed_structured_messages) + self._write_back_structured_messages(data, structured_messages or (), guardrailed_structured_messages) else: # Step 3: Map guardrail responses back to original input structure await self._apply_guardrail_responses_to_input( messages=input_data, - responses=guardrailed_inputs.get("texts", []), + responses=guardrailed_inputs.get("texts", []), # mutable-ok: callee signature takes a list task_mappings=task_mappings, ) @@ -231,10 +346,16 @@ class OpenAIResponsesHandler(BaseTranslation): return data @staticmethod - def _write_back_structured_messages(data: dict, structured_messages: Sequence[AllMessageValues]) -> None: + def _write_back_structured_messages( + data: dict, + original_messages: Sequence[object], + structured_messages: Sequence[AllMessageValues], + ) -> None: + if _patch_rewritten_rows_into_input(data, original_messages, structured_messages): + return input_items, instructions = ( LiteLLMResponsesTransformationHandler().convert_chat_completion_messages_to_responses_api( - list(structured_messages) + list(structured_messages) # mutable-ok: converter signature takes a list ) ) data["input"] = input_items diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 8b1eeb30306..f32599e8098 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -557,6 +557,12 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): hidden_params: Final = getattr(chunk, "_hidden_params", None) if hidden_params is not None: chunk_dict["_hidden_params"] = dict(hidden_params) if isinstance(hidden_params, dict) else hidden_params + if ( + chunk_dict.get("usage") is None + and isinstance(hidden_params, dict) + and hidden_params.get("usage") is not None + ): + chunk_dict["usage"] = hidden_params["usage"] return chunk_dict def create_reasoning_summary_text_done_event( diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index 63f8babf760..87449a4b1e8 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -1253,6 +1253,43 @@ class StructuredRewriteGuardrail(CustomGuardrail): return {**inputs, "structured_messages": rewritten} +class ToolOutputRewriteGuardrail(CustomGuardrail): + """Guardrail that compresses the first tool-result row, the way Headroom does.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + messages = list(inputs.get("structured_messages") or []) + first_tool = next(i for i, m in enumerate(messages) if isinstance(m, dict) and m.get("role") == "tool") + rewritten = [ + {**m, "content": COMPRESSED_MARKER} if i == first_tool else m for i, m in enumerate(messages) + ] + return {**inputs, "structured_messages": rewritten} + + +class DroppingRewriteGuardrail(CustomGuardrail): + """Guardrail that rewrites the first user row and drops the last row, so the + rewrite can only land through the full-conversion fallback.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + messages = list(inputs.get("structured_messages") or []) + first_user = next(i for i, m in enumerate(messages) if isinstance(m, dict) and m.get("role") == "user") + rewritten = [ + {**m, "content": COMPRESSED_MARKER} if i == first_user else m for i, m in enumerate(messages) + ] + return {**inputs, "structured_messages": rewritten[:-1]} + + def _texts(item: dict) -> list[str]: content = item.get("content") if isinstance(content, str): @@ -1296,7 +1333,93 @@ class TestStructuredMessagesWriteBack: assert "instructions" not in result @pytest.mark.asyncio - async def test_developer_item_survives_write_back_as_input_text(self): + async def test_developer_item_preserved_verbatim_by_row_patch(self): + handler = OpenAIResponsesHandler() + developer_item = {"role": "developer", "content": "Always answer in French."} + data = { + "model": "gpt-5.6", + "input": [ + developer_item, + {"role": "user", "content": "memo " * 400}, + {"role": "user", "content": "What is the codename?"}, + ], + } + + result = await handler.process_input_messages(data, StructuredRewriteGuardrail()) + + assert result["input"][0] is developer_item + assert developer_item["content"] == "Always answer in French." + assert _texts(result["input"][1]) == [COMPRESSED_MARKER] + assert _texts(result["input"][2]) == ["What is the codename?"] + + @pytest.mark.asyncio + async def test_reasoning_and_function_call_items_survive_tool_output_compression(self): + handler = OpenAIResponsesHandler() + reasoning_item = { + "id": "rs_123", + "type": "reasoning", + "summary": [], + "encrypted_content": "gAAAAA-signed-reasoning", + } + function_call_item = { + "id": "fc_123", + "type": "function_call", + "call_id": "call_abc", + "name": "read_document", + "arguments": '{"path": "memo.txt"}', + "status": "completed", + } + data = { + "model": "gpt-5.6", + "instructions": "Answer from the memo only.", + "input": [ + reasoning_item, + function_call_item, + {"type": "function_call_output", "call_id": "call_abc", "output": "memo " * 400}, + {"role": "user", "content": "What is the codename?"}, + ], + } + + result = await handler.process_input_messages(data, ToolOutputRewriteGuardrail()) + + assert result["instructions"] == "Answer from the memo only." + assert result["input"][0] is reasoning_item + assert reasoning_item["encrypted_content"] == "gAAAAA-signed-reasoning" + assert result["input"][1] is function_call_item + assert function_call_item["id"] == "fc_123" + assert result["input"][2] == { + "type": "function_call_output", + "call_id": "call_abc", + "output": COMPRESSED_MARKER, + } + assert result["input"][3] == {"role": "user", "content": "What is the codename?"} + + @pytest.mark.asyncio + async def test_web_search_call_item_preserved_verbatim(self): + handler = OpenAIResponsesHandler() + web_search_item = { + "id": "ws_123", + "type": "web_search_call", + "status": "completed", + "action": {"type": "search", "query": "codename memo"}, + } + data = { + "model": "gpt-5.6", + "input": [ + web_search_item, + {"role": "user", "content": "memo " * 400}, + {"role": "user", "content": "What is the codename?"}, + ], + } + + result = await handler.process_input_messages(data, StructuredRewriteGuardrail()) + + assert result["input"][0] is web_search_item + assert _texts(result["input"][1]) == [COMPRESSED_MARKER] + assert _texts(result["input"][2]) == ["What is the codename?"] + + @pytest.mark.asyncio + async def test_row_count_change_falls_back_to_full_conversion(self): handler = OpenAIResponsesHandler() data = { "model": "gpt-5.6", @@ -1307,10 +1430,12 @@ class TestStructuredMessagesWriteBack: ], } - result = await handler.process_input_messages(data, StructuredRewriteGuardrail()) + result = await handler.process_input_messages(data, DroppingRewriteGuardrail()) + assert len(result["input"]) == 2 developer = next(item for item in result["input"] if item.get("role") == "developer") assert developer["content"] == [{"type": "input_text", "text": "Always answer in French."}] + assert _texts(next(item for item in result["input"] if item.get("role") == "user")) == [COMPRESSED_MARKER] @pytest.mark.asyncio async def test_same_inputs_object_back_keeps_the_text_mapping(self): diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py index c297f593617..22c4f46e3fb 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -979,7 +979,7 @@ async def test_responses_request_sends_compressed_input_and_retrieve_tool_upstre result = await OpenAIResponsesHandler().process_input_messages(data=data, guardrail_to_apply=guardrail) assert result["instructions"] == ORIGINAL_MESSAGES[0]["content"] - assert [item["content"][0]["text"] for item in result["input"]] == [ + assert [item["content"] for item in result["input"]] == [ COMPRESSED_MESSAGES_WITH_HASH[0]["content"], ORIGINAL_MESSAGES[2]["content"], ORIGINAL_MESSAGES[3]["content"], diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 823f656ddc5..3ffd44b2fbe 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -24,6 +24,7 @@ from litellm.types.utils import ( ModelResponse, ModelResponseStream, StreamingChoices, + Usage, ) CHAT_COMPLETION_ID = "chatcmpl-77d33d09-effa-4cd2-9c0d-c742d4358256" @@ -523,3 +524,17 @@ async def test_streaming_response_id_falls_back_when_upstream_yields_nothing(): assert response_ids assert len(set(response_ids)) == 1 assert response_ids[0].startswith("resp_") + + +def test_completed_event_restores_usage_hidden_by_stream_options_none(): + final_chunk = _chunk("", finish_reason="stop") + final_chunk._hidden_params = {"usage": Usage(prompt_tokens=117, completion_tokens=5, total_tokens=122)} + iterator = _build_iterator([_chunk("the document says hello"), final_chunk]) + + events = list(iterator) + + completed = next( + event for event in events if getattr(event, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + ) + assert completed.response.usage.input_tokens == 117 + assert completed.response.usage.output_tokens == 5