scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 182

fix: fix tool call loop for compression on streaming /v1/messages

Browse source
This commit is contained in:
Krrish Dholakia 2026-04-15 11:28:18 -07:00
parent df0fdb7bd7
commit e23707dbc4
5 changed files with 1293 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

364
litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py Normal file
View file

@ -0,0 +1,364 @@
"""
Agentic Streaming Iterator for Anthropic Messages
Wraps the raw SSE byte stream from the Anthropic pass-through endpoint,
yields every chunk to the caller (preserving real streaming), collects
all bytes, and on stream exhaustion rebuilds the full Anthropic response
to run through agentic completion hooks. If an agentic hook fires, the
follow-up response is chained as Phase 2 of the same iterator.
"""
import json
from typing import Any, AsyncIterator, Dict, List, Optional, cast
from litellm._logging import verbose_logger
# ---------------------------------------------------------------------------
# SSE parsing helpers (module-level to keep the class lean)
# ---------------------------------------------------------------------------
def _parse_sse_events(raw: bytes) -> List[tuple]:
"""Return a list of (event_type, parsed_data_dict) from raw SSE bytes."""
text = raw.decode("utf-8", errors="replace")
lines = text.split("\n")
events: List[tuple] = []
current_event_type: Optional[str] = None
for line in lines:
stripped = line.strip()
if stripped.startswith("event:"):
current_event_type = stripped[len("event:") :].strip()
continue
if not stripped.startswith("data:"):
continue
data_str = stripped[len("data:") :].strip()
try:
data = json.loads(data_str)
except (json.JSONDecodeError, ValueError):
continue
event_type = current_event_type or data.get("type", "")
current_event_type = None
events.append((event_type, data))
return events
def _handle_message_start(data: Dict, response: Dict) -> None:
msg = data.get("message", {})
response["id"] = msg.get("id", response["id"])
response["model"] = msg.get("model", response["model"])
response["role"] = msg.get("role", response["role"])
usage = msg.get("usage", {})
if usage:
response["usage"]["input_tokens"] = usage.get("input_tokens", 0)
for key in ("cache_creation_input_tokens", "cache_read_input_tokens"):
if key in usage:
response["usage"][key] = usage[key]
def _handle_content_block_start(data: Dict, content_blocks: Dict[int, Dict]) -> None:
idx = data.get("index", len(content_blocks))
block = data.get("content_block", {})
block_type = block.get("type", "text")
_BLOCK_TEMPLATES: Dict[str, Dict] = {
"text": {"type": "text", "text": ""},
"thinking": {"type": "thinking", "thinking": "", "signature": ""},
"redacted_thinking": {
"type": "redacted_thinking",
"data": block.get("data", ""),
},
}
if block_type == "tool_use":
content_blocks[idx] = {
"type": "tool_use",
"id": block.get("id", ""),
"name": block.get("name", ""),
"input": {},
"_partial_json": "",
}
elif block_type in _BLOCK_TEMPLATES:
content_blocks[idx] = dict(_BLOCK_TEMPLATES[block_type])
else:
content_blocks[idx] = dict(block)
def _handle_content_block_delta(data: Dict, content_blocks: Dict[int, Dict]) -> None:
idx = data.get("index", 0)
delta = data.get("delta", {})
delta_type = delta.get("type", "")
block = content_blocks.get(idx)
if block is None:
return
if delta_type == "text_delta":
block["text"] = block.get("text", "") + delta.get("text", "")
elif delta_type == "input_json_delta":
block["_partial_json"] = block.get("_partial_json", "") + delta.get(
"partial_json", ""
)
elif delta_type == "thinking_delta":
block["thinking"] = block.get("thinking", "") + delta.get("thinking", "")
elif delta_type == "signature_delta":
block["signature"] = delta.get("signature", block.get("signature", ""))
def _handle_content_block_stop(data: Dict, content_blocks: Dict[int, Dict]) -> None:
idx = data.get("index", 0)
block = content_blocks.get(idx)
if block and block.get("type") == "tool_use":
partial = block.pop("_partial_json", "")
if partial:
try:
block["input"] = json.loads(partial)
except (json.JSONDecodeError, ValueError):
block["input"] = {"_raw": partial}
def _handle_message_delta(data: Dict, response: Dict) -> None:
delta = data.get("delta", {})
if "stop_reason" in delta:
response["stop_reason"] = delta["stop_reason"]
if "stop_sequence" in delta:
response["stop_sequence"] = delta["stop_sequence"]
usage = data.get("usage", {})
if usage.get("output_tokens") is not None:
response["usage"]["output_tokens"] = usage["output_tokens"]
for key in (
"input_tokens",
"cache_creation_input_tokens",
"cache_read_input_tokens",
):
if key in usage:
response["usage"][key] = usage[key]
class AgenticAnthropicStreamingIterator:
"""
Two-phase async iterator that enables agentic hooks on streaming
Anthropic Messages pass-through responses.
Phase 1: Yield raw SSE bytes from the upstream response while
accumulating them. When the inner iterator is exhausted,
rebuild the full Anthropic response dict and call agentic hooks.
Phase 2: If an agentic hook fires and returns a follow-up response
(streaming or non-streaming), yield those bytes to the caller.
"""
def __init__(
self,
completion_stream: AsyncIterator,
http_handler: Any,
model: str,
messages: List[Dict],
anthropic_messages_provider_config: Any,
anthropic_messages_optional_request_params: Dict,
logging_obj: Any,
custom_llm_provider: str,
kwargs: Dict,
):
self._inner = completion_stream.__aiter__()
self._http_handler = http_handler
self._model = model
self._messages = messages
self._anthropic_messages_provider_config = anthropic_messages_provider_config
self._anthropic_messages_optional_request_params = (
anthropic_messages_optional_request_params
)
self._logging_obj = logging_obj
self._custom_llm_provider = custom_llm_provider
self._kwargs = kwargs
self._collected_bytes: List[bytes] = []
self._stream_exhausted = False
self._hook_processing_done = False
self._follow_up_iterator: Optional[AsyncIterator] = None
def __aiter__(self):
return self
async def __anext__(self) -> bytes:
# Phase 1: yield from upstream, collect bytes
if not self._stream_exhausted:
try:
chunk = await self._inner.__anext__()
self._collected_bytes.append(chunk)
print(
f"\n[AgenticStreaming] Phase 1 — yielding chunk "
f"#{len(self._collected_bytes)} ({len(chunk)} bytes)"
)
return chunk
except StopAsyncIteration:
self._stream_exhausted = True
print(
f"\n[AgenticStreaming] Phase 1 complete — "
f"collected {len(self._collected_bytes)} chunks, "
f"total {sum(len(c) for c in self._collected_bytes)} bytes. "
f"Running agentic hooks..."
)
await self._process_agentic_hooks()
# Fall through to Phase 2
# Phase 2: yield from follow-up stream if one was created
if self._follow_up_iterator is not None:
chunk = await self._follow_up_iterator.__anext__()
print(
f"\n[AgenticStreaming] Phase 2 — yielding follow-up chunk "
f"({len(chunk)} bytes)"
)
return chunk
print("\n[AgenticStreaming] Stream fully exhausted (no follow-up)")
raise StopAsyncIteration
async def _process_agentic_hooks(self) -> None:
"""Rebuild the Anthropic response from collected SSE bytes and call hooks."""
if self._hook_processing_done:
return
self._hook_processing_done = True
if not self._collected_bytes:
print("[AgenticStreaming] No bytes collected, skipping hooks")
return
try:
rebuilt = self._rebuild_anthropic_response_from_sse(self._collected_bytes)
if rebuilt is None:
print("[AgenticStreaming] Could not rebuild response from SSE bytes")
verbose_logger.debug(
"AgenticStreamingIterator: Could not rebuild response from SSE bytes"
)
return
print(
f"[AgenticStreaming] Rebuilt response: id={rebuilt.get('id')}, "
f"model={rebuilt.get('model')}, "
f"stop_reason={rebuilt.get('stop_reason')}, "
f"content_blocks={len(rebuilt.get('content', []))}"
)
content_types = [
f"{b.get('type')}({b.get('name', '')})"
if b.get("type") == "tool_use"
else b.get("type")
for b in rebuilt.get("content", [])
]
print(f"[AgenticStreaming] Content block types: {content_types}")
result = await self._http_handler._call_agentic_completion_hooks(
response=rebuilt,
model=self._model,
messages=self._messages,
anthropic_messages_provider_config=self._anthropic_messages_provider_config,
anthropic_messages_optional_request_params=self._anthropic_messages_optional_request_params,
logging_obj=self._logging_obj,
stream=True,
custom_llm_provider=self._custom_llm_provider,
kwargs=self._kwargs,
)
if result is None:
print(
"[AgenticStreaming] Hooks returned None — "
"no managed tool detected, no follow-up"
)
return
if hasattr(result, "__aiter__"):
print(
f"[AgenticStreaming] Hooks returned async iterator "
f"({type(result).__name__}) — chaining as Phase 2"
)
self._follow_up_iterator = result.__aiter__()
elif isinstance(result, dict):
print(
f"[AgenticStreaming] Hooks returned dict response "
f"(id={result.get('id')}) — wrapping in FakeStreamIterator for Phase 2"
)
from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import (
FakeAnthropicMessagesStreamIterator,
)
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)
fake = FakeAnthropicMessagesStreamIterator(
response=cast(AnthropicMessagesResponse, result)
)
self._follow_up_iterator = fake.__aiter__()
else:
print(
f"[AgenticStreaming] Hooks returned unexpected type: "
f"{type(result).__name__}"
)
verbose_logger.warning(
"AgenticStreamingIterator: Unexpected result type from hooks: %s",
type(result).__name__,
)
except Exception as e:
_call_id = getattr(self._logging_obj, "litellm_call_id", "unknown")
print(
f"[AgenticStreaming] ERROR in hook processing: {e} "
f"(call_id={_call_id})"
)
verbose_logger.exception(
"AgenticStreamingIterator: Error in agentic hook processing "
"[call_id=%s model=%s]: %s",
_call_id,
self._model,
str(e),
)
@staticmethod
def _rebuild_anthropic_response_from_sse(
raw_bytes: List[bytes],
) -> Optional[Dict[str, Any]]:
"""
Parse collected SSE bytes into an Anthropic Messages response dict.
Processes SSE events in order:
- message_start -> envelope (id, model, role, usage)
- content_block_start -> new content block
- content_block_delta -> accumulate text/json/thinking deltas
- content_block_stop -> finalize block
- message_delta -> stop_reason, output usage
- message_stop -> end
"""
events = _parse_sse_events(b"".join(raw_bytes))
response: Dict[str, Any] = {
"id": "",
"type": "message",
"role": "assistant",
"model": "",
"content": [],
"stop_reason": None,
"stop_sequence": None,
"usage": {"input_tokens": 0, "output_tokens": 0},
}
content_blocks: Dict[int, Dict[str, Any]] = {}
saw_message_start = False
for event_type, data in events:
if event_type == "message_start":
saw_message_start = True
_handle_message_start(data, response)
elif event_type == "content_block_start":
_handle_content_block_start(data, content_blocks)
elif event_type == "content_block_delta":
_handle_content_block_delta(data, content_blocks)
elif event_type == "content_block_stop":
_handle_content_block_stop(data, content_blocks)
elif event_type == "message_delta":
_handle_message_delta(data, response)
if not saw_message_start:
return None
for idx in sorted(content_blocks.keys()):
block = content_blocks[idx]
block.pop("_partial_json", None)
response["content"].append(block)
return response

48
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -1978,13 +1978,33 @@ class BaseLLMHTTPHandler:
initial_response: Union[AsyncIterator, AnthropicMessagesResponse]
if stream:
print(
f"\n[llm_http_handler] stream=True for model={model}, "
f"wrapping in AgenticAnthropicStreamingIterator"
)
completion_stream = anthropic_messages_provider_config.get_async_streaming_response_iterator(
model=model,
httpx_response=response,
request_body=request_body,
litellm_logging_obj=logging_obj,
)
initial_response = completion_stream
from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import (
AgenticAnthropicStreamingIterator,
)
initial_response = AgenticAnthropicStreamingIterator(
completion_stream=completion_stream,
http_handler=self,
model=model,
messages=messages,
anthropic_messages_provider_config=anthropic_messages_provider_config,
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
kwargs=kwargs,
)
return initial_response
else:
initial_response = anthropic_messages_provider_config.transform_anthropic_messages_response(
model=model,
@ -1992,7 +2012,7 @@ class BaseLLMHTTPHandler:
logging_obj=logging_obj,
)
# Call agentic completion hooks
# Call agentic completion hooks (non-streaming path only)
final_response = await self._call_agentic_completion_hooks(
response=initial_response,
model=model,
@ -2000,7 +2020,7 @@ class BaseLLMHTTPHandler:
anthropic_messages_provider_config=anthropic_messages_provider_config,
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
logging_obj=logging_obj,
stream=stream or False,
stream=False,
custom_llm_provider=custom_llm_provider,
kwargs=kwargs,
)
@ -4479,7 +4499,12 @@ class BaseLLMHTTPHandler:
max_loops: int,
fingerprints: List[str],
fingerprint: str,
stream: bool = False,
) -> Any:
print(
f"\n[_execute_anthropic_agentic_plan] "
f"model={model}, stream={stream}, depth={depth}/{max_loops}"
)
from litellm.anthropic_interface import messages as anthropic_messages
patch = plan.request_patch or AgenticLoopRequestPatch()
@ -4520,13 +4545,12 @@ class BaseLLMHTTPHandler:
kwargs_for_followup["max_agentic_loops"] = max_loops
kwargs_for_followup["_agentic_loop_fingerprints"] = fingerprints + [fingerprint]
print("optional_params", optional_params)
print("kwargs_for_followup", kwargs_for_followup)
return await anthropic_messages.acreate(
**{
"max_tokens": max_tokens,
"messages": patch.messages,
"model": patch.model or full_model_name,
"stream": stream,
**optional_params,
**kwargs_for_followup,
}
@ -4613,6 +4637,13 @@ class BaseLLMHTTPHandler:
tools = anthropic_messages_optional_request_params.get("tools", [])
depth, max_loops, fingerprints = self._get_agentic_loop_settings(kwargs=kwargs)
print(
f"\n[_call_agentic_completion_hooks] model={model}, stream={stream}, "
f"depth={depth}/{max_loops}, "
f"response_type={type(response).__name__}, "
f"num_callbacks={len(callbacks)}"
)
for callback in callbacks:
try:
if isinstance(callback, CustomLogger):
@ -4630,6 +4661,12 @@ class BaseLLMHTTPHandler:
kwargs=kwargs,
)
print(
f"[_call_agentic_completion_hooks] "
f"callback={callback.__class__.__name__}, "
f"should_run={should_run}"
)
if should_run:
fingerprint = self._fingerprint_agentic_tools(tool_calls)
if fingerprint in fingerprints:
@ -4697,6 +4734,7 @@ class BaseLLMHTTPHandler:
max_loops=max_loops,
fingerprints=fingerprints,
fingerprint=fingerprint,
stream=stream,
)
except Exception as e:

2
litellm/proxy/_new_secret_config.yaml
View file

@ -35,7 +35,7 @@ litellm_settings:
callbacks: ["compression_interception"]
compression_interception_params:
enabled: true
compression_trigger: 1000
compression_trigger: 100000
# # optional:
# # embedding_model: "text-embedding-3-small"
# # embedding_model_params:

93
tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py
View file

@ -180,6 +180,99 @@ async def test_build_agentic_loop_plan_returns_request_patch():
assert "max_tokens" not in plan.request_patch.optional_params
@pytest.mark.asyncio
async def test_should_run_agentic_loop_with_custom_type_tools():
"""Test that async_should_run_agentic_loop returns True when tools contain
litellm_content_retrieve as a custom-typed tool (e.g. Claude Code tool list)
and the model response includes a matching tool_use block."""
logger = CompressionInterceptionLogger()
# Exact tools payload produced by Claude Code – litellm_content_retrieve is
# the final entry and uses type="custom" (not type="function").
tools = [
{
"name": "Agent",
"description": "Launch a new agent to handle complex, multi-step tasks.",
"input_schema": {
"$schema": "https://json-schema.org/draft/2020-12/schema",
"type": "object",
"properties": {
"description": {"type": "string"},
"prompt": {"type": "string"},
},
"required": ["description", "prompt"],
"additionalProperties": False,
},
},
{
"name": "AskUserQuestion",
"description": "Use this tool when you need to ask the user questions.",
"input_schema": {
"$schema": "https://json-schema.org/draft/2020-12/schema",
"type": "object",
"properties": {
"questions": {"type": "array", "items": {"type": "object"}},
},
"required": ["questions"],
"additionalProperties": False,
},
},
{
"name": "Bash",
"description": "Executes a given bash command and returns its output.",
"input_schema": {
"$schema": "https://json-schema.org/draft/2020-12/schema",
"type": "object",
"properties": {"command": {"type": "string"}},
"required": ["command"],
"additionalProperties": False,
},
},
{
"name": "litellm_content_retrieve",
"description": "Retrieve the full content of a file or message that was compressed to save tokens.",
"input_schema": {
"type": "object",
"properties": {
"key": {
"type": "string",
"description": "The identifier of the content to retrieve",
"enum": ["message_0", "HA_UPTIME_ROUTER_SPEC.md", "message_159", "message_160"],
}
},
"required": ["key"],
},
"type": "custom",
},
]
response = {
"content": [
{
"type": "tool_use",
"id": "toolu_abc",
"name": "litellm_content_retrieve",
"input": {"key": "message_0"},
}
]
}
should_run, tools_dict = await logger.async_should_run_agentic_loop(
response=response,
model="claude-3-5-sonnet",
messages=[],
tools=tools,
stream=False,
custom_llm_provider="anthropic",
kwargs={},
)
assert should_run is True
assert tools_dict["tool_type"] == "compression_retrieval"
assert len(tools_dict["tool_calls"]) == 1
assert tools_dict["tool_calls"][0]["input"]["key"] == "message_0"
@pytest.mark.asyncio
async def test_build_agentic_loop_plan_missing_key_fallback():
"""Missing cache keys should produce deterministic fallback content."""

792
tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py Normal file
View file

@ -0,0 +1,792 @@
"""
Tests for AgenticAnthropicStreamingIterator and SSE rebuild helpers.
"""
import json
import os
import sys
from typing import Any, Dict, List, Optional, Tuple
from unittest.mock import AsyncMock, MagicMock
import pytest
sys.path.insert(0, os.path.abspath("../../../../.."))
from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import (
AgenticAnthropicStreamingIterator,
_handle_content_block_delta,
_handle_content_block_start,
_handle_content_block_stop,
_handle_message_delta,
_handle_message_start,
_parse_sse_events,
)
# ---------------------------------------------------------------------------
# Helpers to build SSE byte payloads
# ---------------------------------------------------------------------------
def _sse_event(event_type: str, data: dict) -> bytes:
return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode()
def _build_simple_text_stream() -> List[bytes]:
"""Produce SSE bytes for a simple text response (no tool calls)."""
chunks = []
chunks.append(
_sse_event(
"message_start",
{
"type": "message_start",
"message": {
"id": "msg_123",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-20250514",
"content": [],
"stop_reason": None,
"stop_sequence": None,
"usage": {"input_tokens": 10, "output_tokens": 0},
},
},
)
)
chunks.append(
_sse_event(
"content_block_start",
{
"type": "content_block_start",
"index": 0,
"content_block": {"type": "text", "text": ""},
},
)
)
chunks.append(
_sse_event(
"content_block_delta",
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "text_delta", "text": "Hello, world!"},
},
)
)
chunks.append(
_sse_event("content_block_stop", {"type": "content_block_stop", "index": 0})
)
chunks.append(
_sse_event(
"message_delta",
{
"type": "message_delta",
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
"usage": {"output_tokens": 5},
},
)
)
chunks.append(_sse_event("message_stop", {"type": "message_stop"}))
return chunks
def _build_tool_use_stream() -> List[bytes]:
"""Produce SSE bytes for a response with a tool_use block."""
chunks = []
chunks.append(
_sse_event(
"message_start",
{
"type": "message_start",
"message": {
"id": "msg_tool_456",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-20250514",
"content": [],
"stop_reason": None,
"usage": {"input_tokens": 50, "output_tokens": 0},
},
},
)
)
# thinking block
chunks.append(
_sse_event(
"content_block_start",
{
"type": "content_block_start",
"index": 0,
"content_block": {
"type": "thinking",
"thinking": "",
"signature": "",
},
},
)
)
chunks.append(
_sse_event(
"content_block_delta",
{
"type": "content_block_delta",
"index": 0,
"delta": {
"type": "thinking_delta",
"thinking": "I need to retrieve...",
},
},
)
)
chunks.append(
_sse_event(
"content_block_delta",
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "signature_delta", "signature": "sig_abc"},
},
)
)
chunks.append(
_sse_event("content_block_stop", {"type": "content_block_stop", "index": 0})
)
# tool_use block
chunks.append(
_sse_event(
"content_block_start",
{
"type": "content_block_start",
"index": 1,
"content_block": {
"type": "tool_use",
"id": "toolu_001",
"name": "litellm_content_retrieve",
"input": {},
},
},
)
)
chunks.append(
_sse_event(
"content_block_delta",
{
"type": "content_block_delta",
"index": 1,
"delta": {
"type": "input_json_delta",
"partial_json": '{"key": "section_',
},
},
)
)
chunks.append(
_sse_event(
"content_block_delta",
{
"type": "content_block_delta",
"index": 1,
"delta": {"type": "input_json_delta", "partial_json": '1"}'},
},
)
)
chunks.append(
_sse_event("content_block_stop", {"type": "content_block_stop", "index": 1})
)
chunks.append(
_sse_event(
"message_delta",
{
"type": "message_delta",
"delta": {"stop_reason": "tool_use"},
"usage": {"output_tokens": 20},
},
)
)
chunks.append(_sse_event("message_stop", {"type": "message_stop"}))
return chunks
# ---------------------------------------------------------------------------
# Mock async stream
# ---------------------------------------------------------------------------
class MockAsyncStream:
"""Async iterator that yields a list of byte chunks."""
def __init__(self, chunks: List[bytes]):
self._chunks = list(chunks)
self._idx = 0
def __aiter__(self):
return self
async def __anext__(self) -> bytes:
if self._idx >= len(self._chunks):
raise StopAsyncIteration
chunk = self._chunks[self._idx]
self._idx += 1
return chunk
# ---------------------------------------------------------------------------
# Tests for _parse_sse_events
# ---------------------------------------------------------------------------
class TestParseSSEEvents:
def test_should_parse_single_event(self):
raw = _sse_event(
"message_start", {"type": "message_start", "message": {"id": "1"}}
)
events = _parse_sse_events(raw)
assert len(events) == 1
assert events[0][0] == "message_start"
assert events[0][1]["message"]["id"] == "1"
def test_should_parse_multiple_events(self):
raw = b"".join(_build_simple_text_stream())
events = _parse_sse_events(raw)
event_types = [e[0] for e in events]
assert "message_start" in event_types
assert "content_block_start" in event_types
assert "content_block_delta" in event_types
assert "content_block_stop" in event_types
assert "message_delta" in event_types
assert "message_stop" in event_types
def test_should_skip_malformed_json(self):
raw = b"event: message_start\ndata: {invalid json}\n\n"
events = _parse_sse_events(raw)
assert len(events) == 0
def test_should_handle_empty_bytes(self):
events = _parse_sse_events(b"")
assert events == []
# ---------------------------------------------------------------------------
# Tests for _handle_* helpers
# ---------------------------------------------------------------------------
class TestHandleMessageStart:
def test_should_populate_envelope(self):
response: Dict[str, Any] = {
"id": "",
"model": "",
"role": "assistant",
"usage": {"input_tokens": 0, "output_tokens": 0},
}
data = {
"message": {
"id": "msg_abc",
"model": "claude-sonnet-4-20250514",
"role": "assistant",
"usage": {
"input_tokens": 42,
"cache_creation_input_tokens": 100,
},
}
}
_handle_message_start(data, response)
assert response["id"] == "msg_abc"
assert response["model"] == "claude-sonnet-4-20250514"
assert response["usage"]["input_tokens"] == 42
assert response["usage"]["cache_creation_input_tokens"] == 100
class TestHandleContentBlockStart:
def test_should_create_text_block(self):
blocks: Dict[int, Dict] = {}
data = {"index": 0, "content_block": {"type": "text", "text": ""}}
_handle_content_block_start(data, blocks)
assert blocks[0] == {"type": "text", "text": ""}
def test_should_create_tool_use_block(self):
blocks: Dict[int, Dict] = {}
data = {
"index": 1,
"content_block": {
"type": "tool_use",
"id": "toolu_x",
"name": "my_tool",
"input": {},
},
}
_handle_content_block_start(data, blocks)
assert blocks[1]["type"] == "tool_use"
assert blocks[1]["name"] == "my_tool"
assert blocks[1]["_partial_json"] == ""
def test_should_create_thinking_block(self):
blocks: Dict[int, Dict] = {}
data = {
"index": 0,
"content_block": {"type": "thinking", "thinking": "", "signature": ""},
}
_handle_content_block_start(data, blocks)
assert blocks[0]["type"] == "thinking"
class TestHandleContentBlockDelta:
def test_should_accumulate_text(self):
blocks = {0: {"type": "text", "text": "Hello"}}
_handle_content_block_delta(
{"index": 0, "delta": {"type": "text_delta", "text": " World"}},
blocks,
)
assert blocks[0]["text"] == "Hello World"
def test_should_accumulate_json(self):
blocks = {0: {"type": "tool_use", "_partial_json": '{"key":'}}
_handle_content_block_delta(
{
"index": 0,
"delta": {"type": "input_json_delta", "partial_json": '"val"}'},
},
blocks,
)
assert blocks[0]["_partial_json"] == '{"key":"val"}'
def test_should_ignore_missing_block(self):
blocks: Dict[int, Dict] = {}
_handle_content_block_delta(
{"index": 99, "delta": {"type": "text_delta", "text": "x"}},
blocks,
)
assert 99 not in blocks
class TestHandleContentBlockStop:
def test_should_parse_tool_input_json(self):
blocks = {
0: {
"type": "tool_use",
"input": {},
"_partial_json": '{"key": "section_1"}',
}
}
_handle_content_block_stop({"index": 0}, blocks)
assert blocks[0]["input"] == {"key": "section_1"}
assert "_partial_json" not in blocks[0]
def test_should_handle_invalid_json_gracefully(self):
blocks = {
0: {
"type": "tool_use",
"input": {},
"_partial_json": "not valid json",
}
}
_handle_content_block_stop({"index": 0}, blocks)
assert blocks[0]["input"] == {"_raw": "not valid json"}
class TestHandleMessageDelta:
def test_should_set_stop_reason_and_usage(self):
response: Dict[str, Any] = {
"stop_reason": None,
"stop_sequence": None,
"usage": {"input_tokens": 0, "output_tokens": 0},
}
_handle_message_delta(
{
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
"usage": {"output_tokens": 15},
},
response,
)
assert response["stop_reason"] == "end_turn"
assert response["usage"]["output_tokens"] == 15
# ---------------------------------------------------------------------------
# Tests for _rebuild_anthropic_response_from_sse
# ---------------------------------------------------------------------------
class TestRebuildAnthropicResponse:
def test_should_rebuild_simple_text_response(self):
raw_bytes = _build_simple_text_stream()
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(
raw_bytes
)
assert result is not None
assert result["id"] == "msg_123"
assert result["model"] == "claude-sonnet-4-20250514"
assert result["stop_reason"] == "end_turn"
assert len(result["content"]) == 1
assert result["content"][0]["type"] == "text"
assert result["content"][0]["text"] == "Hello, world!"
assert result["usage"]["input_tokens"] == 10
assert result["usage"]["output_tokens"] == 5
def test_should_rebuild_tool_use_response(self):
raw_bytes = _build_tool_use_stream()
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(
raw_bytes
)
assert result is not None
assert result["id"] == "msg_tool_456"
assert result["stop_reason"] == "tool_use"
assert len(result["content"]) == 2
thinking = result["content"][0]
assert thinking["type"] == "thinking"
assert thinking["thinking"] == "I need to retrieve..."
assert thinking["signature"] == "sig_abc"
tool = result["content"][1]
assert tool["type"] == "tool_use"
assert tool["id"] == "toolu_001"
assert tool["name"] == "litellm_content_retrieve"
assert tool["input"] == {"key": "section_1"}
def test_should_return_none_without_message_start(self):
raw_bytes = [
_sse_event(
"content_block_start",
{
"type": "content_block_start",
"index": 0,
"content_block": {"type": "text"},
},
)
]
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(
raw_bytes
)
assert result is None
def test_should_handle_empty_bytes(self):
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(
[]
)
assert result is None
def test_should_handle_multi_event_chunks(self):
"""When multiple SSE events arrive in a single bytes chunk."""
combined = b"".join(_build_simple_text_stream())
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(
[combined]
)
assert result is not None
assert result["content"][0]["text"] == "Hello, world!"
def test_should_preserve_cache_usage_fields(self):
raw_bytes = [
_sse_event(
"message_start",
{
"type": "message_start",
"message": {
"id": "msg_cache",
"model": "claude-sonnet-4-20250514",
"role": "assistant",
"usage": {
"input_tokens": 100,
"cache_creation_input_tokens": 50,
"cache_read_input_tokens": 30,
},
},
},
),
_sse_event(
"message_delta",
{
"type": "message_delta",
"delta": {"stop_reason": "end_turn"},
"usage": {"output_tokens": 10},
},
),
_sse_event("message_stop", {"type": "message_stop"}),
]
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(
raw_bytes
)
assert result is not None
assert result["usage"]["cache_creation_input_tokens"] == 50
assert result["usage"]["cache_read_input_tokens"] == 30
def test_should_handle_redacted_thinking_block(self):
raw_bytes = [
_sse_event(
"message_start",
{
"type": "message_start",
"message": {
"id": "msg_redact",
"model": "claude-sonnet-4-20250514",
"role": "assistant",
"usage": {"input_tokens": 5},
},
},
),
_sse_event(
"content_block_start",
{
"type": "content_block_start",
"index": 0,
"content_block": {"type": "redacted_thinking", "data": "abc123"},
},
),
_sse_event(
"content_block_stop",
{"type": "content_block_stop", "index": 0},
),
_sse_event(
"message_delta",
{
"type": "message_delta",
"delta": {"stop_reason": "end_turn"},
"usage": {"output_tokens": 1},
},
),
_sse_event("message_stop", {"type": "message_stop"}),
]
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(
raw_bytes
)
assert result is not None
assert result["content"][0]["type"] == "redacted_thinking"
# ---------------------------------------------------------------------------
# Tests for AgenticAnthropicStreamingIterator (Phase 1 / Phase 2)
# ---------------------------------------------------------------------------
class TestAgenticStreamingIteratorPhase1:
@pytest.mark.asyncio
async def test_should_yield_all_chunks_when_no_hook_fires(self):
"""When hooks return None, the wrapper should yield all original chunks."""
chunks = _build_simple_text_stream()
mock_stream = MockAsyncStream(chunks)
mock_handler = MagicMock()
mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=None)
iterator = AgenticAnthropicStreamingIterator(
completion_stream=mock_stream,
http_handler=mock_handler,
model="claude-sonnet-4-20250514",
messages=[{"role": "user", "content": "hi"}],
anthropic_messages_provider_config=MagicMock(),
anthropic_messages_optional_request_params={},
logging_obj=MagicMock(),
custom_llm_provider="anthropic",
kwargs={},
)
collected = []
async for chunk in iterator:
collected.append(chunk)
assert len(collected) == len(chunks)
for orig, got in zip(chunks, collected):
assert orig == got
mock_handler._call_agentic_completion_hooks.assert_awaited_once()
@pytest.mark.asyncio
async def test_should_pass_rebuilt_response_to_hooks(self):
"""The rebuilt dict passed to hooks should match the original stream content."""
chunks = _build_tool_use_stream()
mock_stream = MockAsyncStream(chunks)
captured_response = {}
async def mock_hooks(**kwargs):
captured_response.update(kwargs["response"])
return None
mock_handler = MagicMock()
mock_handler._call_agentic_completion_hooks = mock_hooks
iterator = AgenticAnthropicStreamingIterator(
completion_stream=mock_stream,
http_handler=mock_handler,
model="claude-sonnet-4-20250514",
messages=[],
anthropic_messages_provider_config=MagicMock(),
anthropic_messages_optional_request_params={},
logging_obj=MagicMock(),
custom_llm_provider="anthropic",
kwargs={},
)
async for _ in iterator:
pass
assert captured_response["id"] == "msg_tool_456"
assert captured_response["stop_reason"] == "tool_use"
assert captured_response["content"][1]["name"] == "litellm_content_retrieve"
class TestAgenticStreamingIteratorPhase2:
@pytest.mark.asyncio
async def test_should_chain_follow_up_async_iterator(self):
"""When hooks return an async iterator, Phase 2 should yield from it."""
phase1_chunks = _build_simple_text_stream()
phase2_chunks = [b"follow-up-chunk-1", b"follow-up-chunk-2"]
mock_stream = MockAsyncStream(phase1_chunks)
follow_up = MockAsyncStream(phase2_chunks)
mock_handler = MagicMock()
mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=follow_up)
iterator = AgenticAnthropicStreamingIterator(
completion_stream=mock_stream,
http_handler=mock_handler,
model="claude-sonnet-4-20250514",
messages=[],
anthropic_messages_provider_config=MagicMock(),
anthropic_messages_optional_request_params={},
logging_obj=MagicMock(),
custom_llm_provider="anthropic",
kwargs={},
)
collected = []
async for chunk in iterator:
collected.append(chunk)
assert len(collected) == len(phase1_chunks) + len(phase2_chunks)
assert collected[-2:] == phase2_chunks
@pytest.mark.asyncio
async def test_should_convert_dict_response_to_fake_stream(self):
"""When hooks return a dict, it should be wrapped in FakeAnthropicMessagesStreamIterator."""
phase1_chunks = _build_simple_text_stream()
mock_stream = MockAsyncStream(phase1_chunks)
fake_response = {
"id": "msg_followup",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-20250514",
"content": [{"type": "text", "text": "follow-up answer"}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 100, "output_tokens": 20},
}
mock_handler = MagicMock()
mock_handler._call_agentic_completion_hooks = AsyncMock(
return_value=fake_response
)
iterator = AgenticAnthropicStreamingIterator(
completion_stream=mock_stream,
http_handler=mock_handler,
model="claude-sonnet-4-20250514",
messages=[],
anthropic_messages_provider_config=MagicMock(),
anthropic_messages_optional_request_params={},
logging_obj=MagicMock(),
custom_llm_provider="anthropic",
kwargs={},
)
collected = []
async for chunk in iterator:
collected.append(chunk)
# Phase 1 chunks + Phase 2 fake-stream chunks
assert len(collected) > len(phase1_chunks)
# The follow-up chunks should contain the text from the dict response
phase2_bytes = b"".join(collected[len(phase1_chunks) :])
assert b"follow-up answer" in phase2_bytes
class TestAgenticStreamingIteratorErrorHandling:
@pytest.mark.asyncio
async def test_should_swallow_hook_errors(self):
"""Errors in hook processing should be swallowed; Phase 1 chunks are still yielded."""
chunks = _build_simple_text_stream()
mock_stream = MockAsyncStream(chunks)
mock_handler = MagicMock()
mock_handler._call_agentic_completion_hooks = AsyncMock(
side_effect=RuntimeError("hook exploded")
)
mock_logging = MagicMock()
mock_logging.litellm_call_id = "test_call_123"
iterator = AgenticAnthropicStreamingIterator(
completion_stream=mock_stream,
http_handler=mock_handler,
model="claude-sonnet-4-20250514",
messages=[],
anthropic_messages_provider_config=MagicMock(),
anthropic_messages_optional_request_params={},
logging_obj=mock_logging,
custom_llm_provider="anthropic",
kwargs={},
)
collected = []
async for chunk in iterator:
collected.append(chunk)
# All Phase 1 chunks should still have been yielded
assert len(collected) == len(chunks)
@pytest.mark.asyncio
async def test_should_handle_empty_stream(self):
"""An empty upstream stream should not crash."""
mock_stream = MockAsyncStream([])
mock_handler = MagicMock()
mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=None)
iterator = AgenticAnthropicStreamingIterator(
completion_stream=mock_stream,
http_handler=mock_handler,
model="claude-sonnet-4-20250514",
messages=[],
anthropic_messages_provider_config=MagicMock(),
anthropic_messages_optional_request_params={},
logging_obj=MagicMock(),
custom_llm_provider="anthropic",
kwargs={},
)
collected = []
async for chunk in iterator:
collected.append(chunk)
assert collected == []
# hooks should not be called since no bytes were collected
mock_handler._call_agentic_completion_hooks.assert_not_awaited()
@pytest.mark.asyncio
async def test_should_pass_stream_true_to_hooks(self):
"""The wrapper should always pass stream=True to hooks."""
chunks = _build_simple_text_stream()
mock_stream = MockAsyncStream(chunks)
mock_handler = MagicMock()
mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=None)
iterator = AgenticAnthropicStreamingIterator(
completion_stream=mock_stream,
http_handler=mock_handler,
model="claude-sonnet-4-20250514",
messages=[],
anthropic_messages_provider_config=MagicMock(),
anthropic_messages_optional_request_params={},
logging_obj=MagicMock(),
custom_llm_provider="anthropic",
kwargs={},
)
async for _ in iterator:
pass
call_kwargs = mock_handler._call_agentic_completion_hooks.call_args
assert call_kwargs.kwargs["stream"] is True