diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py new file mode 100644 index 00000000000..ada2dfb4566 --- /dev/null +++ b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py @@ -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 diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 2e65d2843ca..8a4451035a8 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -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: diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 10d51c31244..36c90c28559 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -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: diff --git a/tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py b/tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py index 9c46b74c588..d7ce9ea2ab8 100644 --- a/tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py +++ b/tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py @@ -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.""" diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py new file mode 100644 index 00000000000..b9bda07336f --- /dev/null +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py @@ -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