diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index d2ea7e872a2..64fc62b4722 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -4408,6 +4408,37 @@ class BaseLLMHTTPHandler: stream=stream, kwargs=kwargs_with_provider, ) + + # Check if we need to convert agentic loop response to fake stream + # This happens when the original request was streaming but was converted + # to non-streaming for WebSearch interception, and the agentic loop ran + websearch_converted_stream = ( + logging_obj.model_call_details.get("websearch_interception_converted_stream", False) + if logging_obj is not None + else False + ) + + if websearch_converted_stream and agentic_response is not None: + from typing import cast + + from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import ( + FakeAnthropicMessagesStreamIterator, + ) + from litellm.types.llms.anthropic_messages.anthropic_response import ( + AnthropicMessagesResponse, + ) + + verbose_logger.debug( + "WebSearchInterception: Agentic loop completed, converting response to fake stream" + ) + + # Convert the non-streaming response to a fake stream + if isinstance(agentic_response, dict): + fake_stream = FakeAnthropicMessagesStreamIterator( + response=cast(AnthropicMessagesResponse, agentic_response) + ) + return fake_stream + # First hook that runs agentic loop wins return agentic_response diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py index 5abecb46c99..ee6fc78b9e9 100644 --- a/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py @@ -100,3 +100,110 @@ async def test_internal_flags_filtered_from_followup_kwargs(): # Verify regular kwargs are preserved assert kwargs_for_followup["temperature"] == 0.7 assert kwargs_for_followup["max_tokens"] == 1024 + + +def test_fake_stream_iterator_includes_output_tokens(): + """Test that FakeAnthropicMessagesStreamIterator includes output tokens in message_delta event. + + Regression test for GitHub issue #20187 where output tokens showed as 0 when using + websearch_interception with Claude Code because the streaming response conversion + wasn't properly including the output_tokens from the final response. + """ + from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import ( + FakeAnthropicMessagesStreamIterator, + ) + import json + + # Simulate a response from the agentic loop with usage data + test_response = { + "id": "msg_test123", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5-20250929", + "content": [ + { + "type": "text", + "text": "Here are today's breaking news headlines..." + } + ], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": { + "input_tokens": 38196, + "output_tokens": 395 + } + } + + # Create fake stream iterator + fake_stream = FakeAnthropicMessagesStreamIterator(response=test_response) + + # Collect all chunks + chunks = list(fake_stream) + + # Parse chunks to find message_start and message_delta events + message_start_event = None + message_delta_event = None + + for chunk in chunks: + chunk_str = chunk.decode() + if "event: message_start" in chunk_str: + data_line = chunk_str.split("data: ")[1].strip() + message_start_event = json.loads(data_line) + elif "event: message_delta" in chunk_str: + data_line = chunk_str.split("data: ")[1].strip() + message_delta_event = json.loads(data_line) + + # Verify message_start has input_tokens and output_tokens=0 (as per Anthropic spec) + assert message_start_event is not None, "Should have message_start event" + assert message_start_event["message"]["usage"]["input_tokens"] == 38196 + assert message_start_event["message"]["usage"]["output_tokens"] == 0 # Always 0 in message_start + + # Verify message_delta has the actual output_tokens + assert message_delta_event is not None, "Should have message_delta event" + assert message_delta_event["usage"]["output_tokens"] == 395, \ + f"message_delta should have output_tokens=395, got {message_delta_event['usage']['output_tokens']}" + + +def test_fake_stream_iterator_preserves_stop_reason(): + """Test that FakeAnthropicMessagesStreamIterator preserves stop_reason from the original response. + + This is important for the client to know that the response completed successfully. + """ + from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import ( + FakeAnthropicMessagesStreamIterator, + ) + import json + + # Simulate a response with end_turn stop reason + test_response = { + "id": "msg_test123", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5-20250929", + "content": [ + { + "type": "text", + "text": "Response text" + } + ], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": { + "input_tokens": 100, + "output_tokens": 50 + } + } + + fake_stream = FakeAnthropicMessagesStreamIterator(response=test_response) + chunks = list(fake_stream) + + # Find message_delta event + message_delta_event = None + for chunk in chunks: + chunk_str = chunk.decode() + if "event: message_delta" in chunk_str: + data_line = chunk_str.split("data: ")[1].strip() + message_delta_event = json.loads(data_line) + + assert message_delta_event is not None + assert message_delta_event["delta"]["stop_reason"] == "end_turn"