diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 573ba85416f..52a91a1d06d 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -97,9 +97,12 @@ from litellm.types.containers.main import ( ) from litellm.types.files import StreamingMediaUploadConfig, TwoStepFileUploadConfig from litellm.types.integrations.custom_logger import ( + NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES, AgenticLoopPlan, AgenticLoopRequestPatch, AgenticLoopSafetyError, + converted_stream_requested, + is_interception_internal_key, ) from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, @@ -2932,10 +2935,7 @@ class BaseLLMHTTPHandler: ) result: Final = final_response if final_response is not None else initial_response - interception_converted_stream: Final = litellm_params.get( - "_code_interpreter_interception_converted_stream" - ) or litellm_params.get("_websearch_interception_converted_stream") - if interception_converted_stream and not litellm_params.get("_agentic_loop_depth"): + if converted_stream_requested(litellm_params) and not litellm_params.get("_agentic_loop_depth"): return self._wrap_responses_response_as_fake_stream( result=result, model=model, @@ -5399,8 +5399,7 @@ class BaseLLMHTTPHandler: kwargs_for_followup: Final = { k: v for k, v in kwargs.items() - if not k.startswith("_websearch_interception") - and not k.startswith("_compression_interception") + if not is_interception_internal_key(k, prefixes=NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES) and k != "_code_interpreter_interception_converted_stream" and k not in internal_keys and k not in optional_params diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index e2d2fffb2df..08753bdaed9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -49,6 +49,9 @@ if TYPE_CHECKING: from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel BYPASS_HEADER: Final = "x-headroom-bypass" +_STREAM_CONVERTIBLE_CALL_TYPES: Final = frozenset( + (CallTypes.completion, CallTypes.acompletion, CallTypes.responses, CallTypes.aresponses) +) HEADROOM_RETRIEVE_TOOL_NAME: Final = "headroom_retrieve" _HASH_PATTERN: Final = re.compile(r"hash=([a-f0-9]{24})") _HASH_CACHE_TTL_SECONDS: Final = 15 * 60 @@ -724,7 +727,7 @@ class HeadroomGuardrail(CustomGuardrail): ) -> dict[str, Any] | None: # mutable-ok: overrides CustomLogger hook whose contract is a plain dict base_result: Final = await super().async_pre_call_deployment_hook(kwargs, call_type) effective: Final = base_result if base_result is not None else kwargs - if call_type not in (CallTypes.completion, CallTypes.acompletion): + if call_type not in _STREAM_CONVERTIBLE_CALL_TYPES: return base_result if not effective.get("stream"): return base_result diff --git a/litellm/types/integrations/custom_logger.py b/litellm/types/integrations/custom_logger.py index 9a714e1724e..5de58a20242 100644 --- a/litellm/types/integrations/custom_logger.py +++ b/litellm/types/integrations/custom_logger.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Any, Final from pydantic import BaseModel, Field @@ -29,6 +30,13 @@ def is_interception_internal_key( return any(key.startswith(prefix) for prefix in prefixes) +CONVERTED_STREAM_KEYS: Final = frozenset(f"{prefix}_converted_stream" for prefix in INTERCEPTION_INTERNAL_PREFIXES) + + +def converted_stream_requested(params: Mapping[str, object]) -> bool: + return any(bool(params.get(key)) for key in CONVERTED_STREAM_KEYS) + + class AgenticLoopSafetyError(ValueError): """ Raised when an agentic-loop safety rail refuses a rerun. 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 37b9852d3e1..e1c4cc0922a 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -1984,6 +1984,58 @@ def _openai_text_payload(content: str) -> dict: return _openai_completion_payload({"role": "assistant", "content": content}, "stop") +def _responses_retrieve_tool_definition() -> dict: + return {"type": "function", **_retrieve_tool_definition()["function"]} + + +def _openai_responses_payload(output_item: dict) -> dict: + return { + "id": "resp_ccr", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-4o", + "output": [output_item], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "error": None, + "incomplete_details": None, + "instructions": None, + "metadata": {}, + "temperature": 1.0, + "top_p": 1.0, + "text": {"format": {"type": "text"}}, + "truncation": "disabled", + } + + +def _openai_responses_retrieve_call_payload() -> dict: + return _openai_responses_payload( + { + "type": "function_call", + "id": "fc_ccr", + "call_id": "call_ccr", + "name": HEADROOM_RETRIEVE_TOOL_NAME, + "arguments": json.dumps({"hash": CCR_HASH}), + "status": "completed", + } + ) + + +def _openai_responses_text_payload(text: str) -> dict: + return _openai_responses_payload( + { + "type": "message", + "id": "msg_ccr", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": text, "annotations": []}], + } + ) + + @pytest.mark.parametrize( "call_type, stream, tools, expect_conversion", [ @@ -1992,12 +2044,14 @@ def _openai_text_payload(content: str) -> dict: (CallTypes.acompletion, False, [_retrieve_tool_definition()], False), (CallTypes.acompletion, True, [{"type": "function", "function": {"name": "get_weather"}}], False), (CallTypes.acompletion, True, None, False), - (CallTypes.aresponses, True, [_retrieve_tool_definition()], False), + (CallTypes.aresponses, True, [_retrieve_tool_definition()], True), + (CallTypes.responses, True, [_responses_retrieve_tool_definition()], True), + (CallTypes.aresponses, False, [_retrieve_tool_definition()], False), (CallTypes.anthropic_messages, True, [_retrieve_tool_definition()], False), ], ) @pytest.mark.asyncio -async def test_pre_call_deployment_hook_converts_stream_only_for_ccr_chat_completions( +async def test_pre_call_deployment_hook_converts_stream_only_for_ccr_chat_completions_and_responses( guardrail: HeadroomGuardrail, call_type: CallTypes, stream: bool, @@ -2128,6 +2182,65 @@ async def test_streaming_chat_completion_resolves_ccr_retrieval_end_to_end( assert not any(key.startswith("_headroom_interception") for key in followup_body) +@pytest.mark.asyncio +async def test_streaming_responses_resolves_ccr_retrieval_end_to_end( + guardrail: HeadroomGuardrail, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +): + """Regression test for LIT-6481: streaming /v1/responses must resolve the + retrieve tool call server-side exactly like streaming /chat/completions does, + instead of streaming a headroom_retrieve function_call to the client.""" + original_content = "the full uncompressed document" + final_answer = "the document says hello" + guardrail._issued_hashes_by_call_id["ccr-call-id"] = ( + frozenset({CCR_HASH}), + time.monotonic() + 999, + ) + + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + upstream = respx_mock.post("https://api.openai.com/v1/responses").mock( + side_effect=[ + httpx.Response(200, json=_openai_responses_retrieve_call_payload()), + httpx.Response(200, json=_openai_responses_text_payload(final_answer)), + ] + ) + + with patch.object( + guardrail.async_handler, + "get", + new_callable=AsyncMock, + return_value=_make_retrieve_response(original_content), + ) as mock_get: + response = await litellm.aresponses( + model="openai/gpt-4o", + input=[{"role": "user", "content": f"summarize hash={CCR_HASH}"}], + tools=[_responses_retrieve_tool_definition()], + stream=True, + litellm_call_id="ccr-call-id", + ) + events = [event async for event in response] + + streamed_text = "".join( + getattr(event, "delta", "") for event in events if getattr(event, "type", None) == "response.output_text.delta" + ) + assert streamed_text == final_answer + assert not any("function_call" in str(getattr(event, "type", "")) for event in events) + assert not any( + getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events + ) + mock_get.assert_called_once() + assert CCR_HASH in (mock_get.call_args.kwargs.get("url") or mock_get.call_args.args[0]) + + assert len(upstream.calls) == 2 + followup_body = json.loads(upstream.calls[1].request.content) + assert not followup_body.get("stream") + assert original_content in json.dumps(followup_body["input"]) + assert not any(key.startswith("_headroom_interception") for key in followup_body) + + # --------------------------------------------------------------------------- # LIT-5018: the turn the model is being asked to act on is never compressed. #