diff --git a/.sampo/changesets/anthropic-stream-block-index.md b/.sampo/changesets/anthropic-stream-block-index.md new file mode 100644 index 000000000..fc48bd1ce --- /dev/null +++ b/.sampo/changesets/anthropic-stream-block-index.md @@ -0,0 +1,5 @@ +--- +pypi/posthog: patch +--- + +Keep streamed Anthropic tool call arguments when the response also contains server tool blocks, such as web search. diff --git a/posthog/ai/anthropic/_anthropic_stream.py b/posthog/ai/anthropic/_anthropic_stream.py index 4b94913ea..e08f6f65e 100644 --- a/posthog/ai/anthropic/_anthropic_stream.py +++ b/posthog/ai/anthropic/_anthropic_stream.py @@ -21,6 +21,10 @@ def __init__(self) -> None: self.tools_in_progress: Dict[str, ToolInProgress] = {} self.current_text_block: Optional[StreamingContentBlock] = None self.stop_reason: Optional[str] = None + # Delta and stop events address blocks by their index in the message. That + # index also counts blocks content_blocks leaves out (server_tool_use, + # web_search_tool_result, ...), so keep a list that lines up with it. + self._blocks_by_index: List[StreamingContentBlock] = [] def consume(self, event: Any) -> None: event_usage = extract_anthropic_usage_from_event(event) @@ -29,6 +33,8 @@ def consume(self, event: Any) -> None: if getattr(event, "type", None) == "content_block_start": block, tool = handle_anthropic_content_block_start(event) + self._blocks_by_index.append(block or {}) + if block: self.content_blocks.append(block) if block.get("type") in ("text", "thinking"): @@ -45,12 +51,14 @@ def consume(self, event: Any) -> None: if delta_text: self.accumulated_content += delta_text - handle_anthropic_tool_delta(event, self.content_blocks, self.tools_in_progress) + handle_anthropic_tool_delta( + event, self._blocks_by_index, self.tools_in_progress + ) if getattr(event, "type", None) == "content_block_stop": self.current_text_block = None finalize_anthropic_tool_input( - event, self.content_blocks, self.tools_in_progress + event, self._blocks_by_index, self.tools_in_progress ) if getattr(event, "type", None) == "message_delta": diff --git a/posthog/test/ai/anthropic/test_anthropic.py b/posthog/test/ai/anthropic/test_anthropic.py index dc9b2baf3..84275e9ad 100644 --- a/posthog/test/ai/anthropic/test_anthropic.py +++ b/posthog/test/ai/anthropic/test_anthropic.py @@ -1212,6 +1212,88 @@ def test_streaming_with_tool_calls(mock_client, mock_anthropic_stream_with_tools assert "output_tokens" in props["$ai_usage"] +def test_streaming_tool_call_after_server_tool_blocks(mock_client): + """Stream events address content blocks by their index in the message. Blocks the + wrapper does not format (server_tool_use, web_search_tool_result) still take an + index, so a later client tool call must keep its streamed arguments.""" + from anthropic.types import ( + InputJSONDelta, + ServerToolUseBlock, + ToolUseBlock, + WebSearchToolResultBlock, + ) + + def content_block(index, block, deltas=()): + yield RawContentBlockStartEvent( + type="content_block_start", index=index, content_block=block + ) + for delta in deltas: + yield RawContentBlockDeltaEvent( + type="content_block_delta", index=index, delta=delta + ) + yield RawContentBlockStopEvent(type="content_block_stop", index=index) + + def stream_generator(): + yield from content_block( + 0, + ServerToolUseBlock( + type="server_tool_use", id="srvtoolu_1", name="web_search", input={} + ), + [ + InputJSONDelta( + type="input_json_delta", partial_json='{"query": "weather sf"}' + ) + ], + ) + yield from content_block( + 1, + WebSearchToolResultBlock( + type="web_search_tool_result", tool_use_id="srvtoolu_1", content=[] + ), + ) + yield from content_block( + 2, + TextBlock(type="text", text=""), + [TextDelta(type="text_delta", text="Booking a table.")], + ) + yield from content_block( + 3, + ToolUseBlock(type="tool_use", id="toolu_1", name="book_table", input={}), + [ + InputJSONDelta(type="input_json_delta", partial_json='{"city": '), + InputJSONDelta(type="input_json_delta", partial_json='"SF"}'), + ], + ) + + with patch( + "anthropic.resources.Messages.create", + return_value=stream_generator(), + ): + client = Anthropic(api_key="test-key", posthog_client=mock_client) + response = client.messages.create( + model="claude-sonnet-4-20250514", + messages=[{"role": "user", "content": "Book dinner if it's sunny"}], + stream=True, + posthog_distinct_id="test-id", + ) + list(response) + + props = mock_client.capture.call_args[1]["properties"] + assert props["$ai_output_choices"] == [ + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Booking a table."}, + { + "type": "function", + "id": "toolu_1", + "function": {"name": "book_table", "arguments": {"city": "SF"}}, + }, + ], + } + ] + + def test_async_streaming_with_tool_calls(mock_client, mock_anthropic_stream_with_tools): """Test that tool calls are properly captured in async streaming mode.""" import asyncio