mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
fix(websearch_interception): convert agentic loop response to streaming format when original request was streaming
Fixes #20187 - When using websearch_interception in Bedrock with Claude Code: 1. Output tokens were showing as 0 because the agentic loop response wasn't being converted back to streaming format 2. The response from the agentic loop (follow-up request) was returned as a non-streaming dict, but Claude Code expects a streaming response This fix adds streaming format conversion for the agentic loop response when the original request was streaming (detected via the websearch_interception_converted_stream flag in logging_obj). The fix ensures: - Output tokens are correctly included in the message_delta event - stop_reason is properly preserved - The response format matches what Claude Code expects
This commit is contained in:
parent
51af66fdb2
commit
94de0e5403
2 changed files with 138 additions and 0 deletions
|
|
@ -4408,6 +4408,37 @@ class BaseLLMHTTPHandler:
|
||||||
stream=stream,
|
stream=stream,
|
||||||
kwargs=kwargs_with_provider,
|
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
|
# First hook that runs agentic loop wins
|
||||||
return agentic_response
|
return agentic_response
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -100,3 +100,110 @@ async def test_internal_flags_filtered_from_followup_kwargs():
|
||||||
# Verify regular kwargs are preserved
|
# Verify regular kwargs are preserved
|
||||||
assert kwargs_for_followup["temperature"] == 0.7
|
assert kwargs_for_followup["temperature"] == 0.7
|
||||||
assert kwargs_for_followup["max_tokens"] == 1024
|
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"
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue