From 2928f7dac45cff6ca5593b76cf7f6193167ff6fb Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Fri, 21 Aug 2026 16:04:33 +0200 Subject: [PATCH 1/2] refactor(anthropic): share streaming accumulator --- posthog/ai/anthropic/anthropic.py | 86 ++----------- posthog/ai/anthropic/anthropic_async.py | 88 ++----------- posthog/ai/anthropic/anthropic_stream.py | 60 +++++++++ posthog/test/ai/anthropic/test_anthropic.py | 134 ++++++++++++++++++++ 4 files changed, 213 insertions(+), 155 deletions(-) create mode 100644 posthog/ai/anthropic/anthropic_stream.py diff --git a/posthog/ai/anthropic/anthropic.py b/posthog/ai/anthropic/anthropic.py index a25bedc4f..10866f698 100644 --- a/posthog/ai/anthropic/anthropic.py +++ b/posthog/ai/anthropic/anthropic.py @@ -11,18 +11,9 @@ from typing import Any, Dict, List, Optional from ..stream import _StreamWrapper -from posthog.ai.types import StreamingContentBlock, TokenUsage, ToolInProgress -from posthog.ai.utils import ( - call_llm_and_track_usage, - merge_usage_stats, -) -from posthog.ai.anthropic.anthropic_converter import ( - extract_anthropic_usage_from_event, - handle_anthropic_content_block_start, - handle_anthropic_text_delta, - handle_anthropic_tool_delta, - finalize_anthropic_tool_input, -) +from ..types import StreamingContentBlock, TokenUsage +from ..utils import call_llm_and_track_usage +from .anthropic_stream import AnthropicStreamAccumulator from posthog.client import Client as PostHogClient from posthog import setup @@ -180,72 +171,13 @@ def _track_streaming_response( kwargs: Dict[str, Any], start_time: float, ): - usage_stats: TokenUsage = TokenUsage(input_tokens=0, output_tokens=0) - accumulated_content = "" - content_blocks: List[StreamingContentBlock] = [] - tools_in_progress: Dict[str, ToolInProgress] = {} - current_text_block: Optional[StreamingContentBlock] = None - stop_reason: Optional[str] = None + accumulator = AnthropicStreamAccumulator() def generator(): - nonlocal usage_stats - nonlocal accumulated_content - nonlocal content_blocks - nonlocal tools_in_progress - nonlocal current_text_block - nonlocal stop_reason - try: for event in response: - # Extract usage stats from event - event_usage = extract_anthropic_usage_from_event(event) - merge_usage_stats(usage_stats, event_usage) - - # Handle content block start events - if hasattr(event, "type") and event.type == "content_block_start": - block, tool = handle_anthropic_content_block_start(event) - - if block: - content_blocks.append(block) - - if block.get("type") in ("text", "thinking"): - current_text_block = block - else: - current_text_block = None - - if tool: - tool_id = tool["block"].get("id") - if tool_id: - tools_in_progress[tool_id] = tool - - # Handle text delta events - delta_text = handle_anthropic_text_delta(event, current_text_block) - - if delta_text: - accumulated_content += delta_text - - # Handle tool input delta events - handle_anthropic_tool_delta( - event, content_blocks, tools_in_progress - ) - - # Handle content block stop events - if hasattr(event, "type") and event.type == "content_block_stop": - current_text_block = None - finalize_anthropic_tool_input( - event, content_blocks, tools_in_progress - ) - - # Capture stop reason from message_delta events - if hasattr(event, "type") and event.type == "message_delta": - delta = getattr(event, "delta", None) - if delta is not None: - delta_stop_reason = getattr(delta, "stop_reason", None) - if delta_stop_reason is not None: - stop_reason = delta_stop_reason - + accumulator.consume(event) yield event - finally: end_time = time.time() latency = end_time - start_time @@ -257,11 +189,11 @@ def generator(): posthog_privacy_mode, posthog_groups, kwargs, - usage_stats, + accumulator.usage_stats, latency, - content_blocks, - accumulated_content, - stop_reason=stop_reason, + accumulator.content_blocks, + accumulator.accumulated_content, + stop_reason=accumulator.stop_reason, ) return _StreamWrapper(generator(), stream=response) diff --git a/posthog/ai/anthropic/anthropic_async.py b/posthog/ai/anthropic/anthropic_async.py index eb642de8b..3e2bf71cf 100644 --- a/posthog/ai/anthropic/anthropic_async.py +++ b/posthog/ai/anthropic/anthropic_async.py @@ -11,19 +11,10 @@ from typing import Any, Dict, List, Optional from posthog import setup -from posthog.ai.stream import AsyncStreamWrapper -from posthog.ai.types import StreamingContentBlock, TokenUsage, ToolInProgress -from posthog.ai.utils import ( - call_llm_and_track_usage_async, - merge_usage_stats, -) -from posthog.ai.anthropic.anthropic_converter import ( - extract_anthropic_usage_from_event, - handle_anthropic_content_block_start, - handle_anthropic_text_delta, - handle_anthropic_tool_delta, - finalize_anthropic_tool_input, -) +from ..stream import AsyncStreamWrapper +from ..types import StreamingContentBlock, TokenUsage +from ..utils import call_llm_and_track_usage_async +from .anthropic_stream import AnthropicStreamAccumulator from posthog.client import Client as PostHogClient @@ -180,72 +171,13 @@ def _track_streaming_response( kwargs: Dict[str, Any], start_time: float, ): - usage_stats: TokenUsage = TokenUsage(input_tokens=0, output_tokens=0) - accumulated_content = "" - content_blocks: List[StreamingContentBlock] = [] - tools_in_progress: Dict[str, ToolInProgress] = {} - current_text_block: Optional[StreamingContentBlock] = None - stop_reason: Optional[str] = None + accumulator = AnthropicStreamAccumulator() async def generator(): - nonlocal usage_stats - nonlocal accumulated_content - nonlocal content_blocks - nonlocal tools_in_progress - nonlocal current_text_block - nonlocal stop_reason - try: async for event in response: - # Extract usage stats from event - event_usage = extract_anthropic_usage_from_event(event) - merge_usage_stats(usage_stats, event_usage) - - # Handle content block start events - if hasattr(event, "type") and event.type == "content_block_start": - block, tool = handle_anthropic_content_block_start(event) - - if block: - content_blocks.append(block) - - if block.get("type") in ("text", "thinking"): - current_text_block = block - else: - current_text_block = None - - if tool: - tool_id = tool["block"].get("id") - if tool_id: - tools_in_progress[tool_id] = tool - - # Handle text delta events - delta_text = handle_anthropic_text_delta(event, current_text_block) - - if delta_text: - accumulated_content += delta_text - - # Handle tool input delta events - handle_anthropic_tool_delta( - event, content_blocks, tools_in_progress - ) - - # Handle content block stop events - if hasattr(event, "type") and event.type == "content_block_stop": - current_text_block = None - finalize_anthropic_tool_input( - event, content_blocks, tools_in_progress - ) - - # Capture stop reason from message_delta events - if hasattr(event, "type") and event.type == "message_delta": - delta = getattr(event, "delta", None) - if delta is not None: - delta_stop_reason = getattr(delta, "stop_reason", None) - if delta_stop_reason is not None: - stop_reason = delta_stop_reason - + accumulator.consume(event) yield event - finally: end_time = time.time() latency = end_time - start_time @@ -257,11 +189,11 @@ async def generator(): posthog_privacy_mode, posthog_groups, kwargs, - usage_stats, + accumulator.usage_stats, latency, - content_blocks, - accumulated_content, - stop_reason=stop_reason, + accumulator.content_blocks, + accumulator.accumulated_content, + stop_reason=accumulator.stop_reason, ) return AsyncStreamWrapper(generator(), stream=response) diff --git a/posthog/ai/anthropic/anthropic_stream.py b/posthog/ai/anthropic/anthropic_stream.py new file mode 100644 index 000000000..b2dc69c5d --- /dev/null +++ b/posthog/ai/anthropic/anthropic_stream.py @@ -0,0 +1,60 @@ +from typing import Any, Dict, List, Optional + +from ..types import StreamingContentBlock, TokenUsage, ToolInProgress +from ..utils import merge_usage_stats +from .anthropic_converter import ( + extract_anthropic_usage_from_event, + finalize_anthropic_tool_input, + handle_anthropic_content_block_start, + handle_anthropic_text_delta, + handle_anthropic_tool_delta, +) + + +class AnthropicStreamAccumulator: + """Accumulates sync-neutral capture state from Anthropic stream events.""" + + def __init__(self) -> None: + self.usage_stats: TokenUsage = TokenUsage(input_tokens=0, output_tokens=0) + self.accumulated_content = "" + self.content_blocks: List[StreamingContentBlock] = [] + self.tools_in_progress: Dict[str, ToolInProgress] = {} + self.current_text_block: Optional[StreamingContentBlock] = None + self.stop_reason: Optional[str] = None + + def consume(self, event: Any) -> None: + event_usage = extract_anthropic_usage_from_event(event) + merge_usage_stats(self.usage_stats, event_usage) + + if getattr(event, "type", None) == "content_block_start": + block, tool = handle_anthropic_content_block_start(event) + + if block: + self.content_blocks.append(block) + if block.get("type") in ("text", "thinking"): + self.current_text_block = block + else: + self.current_text_block = None + + if tool: + tool_id = tool["block"].get("id") + if tool_id: + self.tools_in_progress[tool_id] = tool + + delta_text = handle_anthropic_text_delta(event, self.current_text_block) + if delta_text: + self.accumulated_content += delta_text + + handle_anthropic_tool_delta(event, self.content_blocks, 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 + ) + + if getattr(event, "type", None) == "message_delta": + delta = getattr(event, "delta", None) + delta_stop_reason = getattr(delta, "stop_reason", None) + if delta_stop_reason is not None: + self.stop_reason = delta_stop_reason diff --git a/posthog/test/ai/anthropic/test_anthropic.py b/posthog/test/ai/anthropic/test_anthropic.py index f7322a637..dc9b2baf3 100644 --- a/posthog/test/ai/anthropic/test_anthropic.py +++ b/posthog/test/ai/anthropic/test_anthropic.py @@ -1779,6 +1779,16 @@ def close(self): self.closed = True +class FailingRecordingStream(RecordingStream): + def __next__(self): + raise RuntimeError("stream failed") + + +class FailingRecordingAsyncStream(RecordingAsyncStream): + async def __anext__(self): + raise RuntimeError("stream failed") + + def _anthropic_stream_events(): final = MockStreamEvent("message_delta") final.usage = MockUsage( @@ -1827,6 +1837,107 @@ def _anthropic_raw_stream_events(): ] +@pytest.mark.asyncio +async def test_streaming_sync_async_accumulation_parity(mock_client): + kwargs = { + "model": "claude-haiku-4-5", + "messages": [{"role": "user", "content": "Foo"}], + "max_tokens": 1, + "stream": True, + } + + sync_client = Anthropic(api_key="test-key", posthog_client=mock_client) + sync_response = sync_client.messages._track_streaming_response( + RecordingStream(_anthropic_raw_stream_events()), + "test-user", + "test-trace", + None, + False, + None, + kwargs, + 0, + ) + list(sync_response) + sync_properties = mock_client.capture.call_args.kwargs["properties"] + + mock_client.capture.reset_mock() + + async_client = AsyncAnthropic(api_key="test-key", posthog_client=mock_client) + async_response = async_client.messages._track_streaming_response( + RecordingAsyncStream(_anthropic_raw_stream_events()), + "test-user", + "test-trace", + None, + False, + None, + kwargs, + 0, + ) + [event async for event in async_response] + async_properties = mock_client.capture.call_args.kwargs["properties"] + + parity_keys = ( + "$ai_input", + "$ai_output_choices", + "$ai_input_tokens", + "$ai_output_tokens", + "$ai_usage", + "$ai_stop_reason", + ) + assert {key: sync_properties[key] for key in parity_keys} == { + key: async_properties[key] for key in parity_keys + } + + +@pytest.mark.asyncio +async def test_streaming_sync_async_exceptions_capture_and_close(mock_client): + kwargs = { + "model": "claude-haiku-4-5", + "messages": [{"role": "user", "content": "Foo"}], + "stream": True, + } + + sync_source = FailingRecordingStream([]) + sync_client = Anthropic(api_key="test-key", posthog_client=mock_client) + sync_response = sync_client.messages._track_streaming_response( + sync_source, + None, + "test-trace", + None, + False, + None, + kwargs, + 0, + ) + with pytest.raises(RuntimeError, match="stream failed"): + with sync_response as stream: + list(stream) + + assert sync_source.closed is True + assert mock_client.capture.call_count == 1 + + mock_client.capture.reset_mock() + + async_source = FailingRecordingAsyncStream([]) + async_client = AsyncAnthropic(api_key="test-key", posthog_client=mock_client) + async_response = async_client.messages._track_streaming_response( + async_source, + None, + "test-trace", + None, + False, + None, + kwargs, + 0, + ) + with pytest.raises(RuntimeError, match="stream failed"): + async with async_response as stream: + [event async for event in stream] + + assert async_source.closed is True + assert mock_client.capture.call_count == 1 + + def test_messages_stream_preserves_native_manager_helpers_close_and_tracking( mock_client, ): @@ -1940,6 +2051,29 @@ async def mock_async_create(**kwargs): assert mock_client.capture.call_count == 1 +def test_messages_streaming_early_exit_closes_provider_stream(mock_client): + source = RecordingStream(_anthropic_stream_events()) + + with patch( + "anthropic.resources.messages.Messages.create", + return_value=source, + ): + client = Anthropic(api_key="test-key", posthog_client=mock_client) + response = client.messages.create( + model="claude-3-opus-20240229", + messages=[{"role": "user", "content": "Foo"}], + stream=True, + max_tokens=1, + ) + + with response as stream: + for _ in stream: + break + + assert source.closed is True + assert mock_client.capture.call_count == 1 + + @pytest.mark.asyncio async def test_async_messages_streaming_early_exit_closes_provider_stream(mock_client): """Breaking out early must close the underlying Anthropic stream and still From 7aca336325d8a08d23b83a6bad886bdb33aa9512 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Sat, 22 Aug 2026 08:53:28 +0200 Subject: [PATCH 2/2] fix(anthropic): preserve public API after stream refactor --- ...thropic_stream.py => _anthropic_stream.py} | 2 +- posthog/ai/anthropic/anthropic.py | 22 +++++++++++++---- posthog/ai/anthropic/anthropic_async.py | 24 +++++++++++++++---- 3 files changed, 38 insertions(+), 10 deletions(-) rename posthog/ai/anthropic/{anthropic_stream.py => _anthropic_stream.py} (98%) diff --git a/posthog/ai/anthropic/anthropic_stream.py b/posthog/ai/anthropic/_anthropic_stream.py similarity index 98% rename from posthog/ai/anthropic/anthropic_stream.py rename to posthog/ai/anthropic/_anthropic_stream.py index b2dc69c5d..3426e1d4b 100644 --- a/posthog/ai/anthropic/anthropic_stream.py +++ b/posthog/ai/anthropic/_anthropic_stream.py @@ -11,7 +11,7 @@ ) -class AnthropicStreamAccumulator: +class _AnthropicStreamAccumulator: """Accumulates sync-neutral capture state from Anthropic stream events.""" def __init__(self) -> None: diff --git a/posthog/ai/anthropic/anthropic.py b/posthog/ai/anthropic/anthropic.py index 10866f698..a474fefcf 100644 --- a/posthog/ai/anthropic/anthropic.py +++ b/posthog/ai/anthropic/anthropic.py @@ -11,9 +11,23 @@ from typing import Any, Dict, List, Optional from ..stream import _StreamWrapper -from ..types import StreamingContentBlock, TokenUsage -from ..utils import call_llm_and_track_usage -from .anthropic_stream import AnthropicStreamAccumulator +from ..types import ( + StreamingContentBlock as StreamingContentBlock, + TokenUsage as TokenUsage, + ToolInProgress as ToolInProgress, +) +from ..utils import ( + call_llm_and_track_usage as call_llm_and_track_usage, + merge_usage_stats as merge_usage_stats, +) +from ._anthropic_stream import _AnthropicStreamAccumulator +from .anthropic_converter import ( + extract_anthropic_usage_from_event as extract_anthropic_usage_from_event, + finalize_anthropic_tool_input as finalize_anthropic_tool_input, + handle_anthropic_content_block_start as handle_anthropic_content_block_start, + handle_anthropic_text_delta as handle_anthropic_text_delta, + handle_anthropic_tool_delta as handle_anthropic_tool_delta, +) from posthog.client import Client as PostHogClient from posthog import setup @@ -171,7 +185,7 @@ def _track_streaming_response( kwargs: Dict[str, Any], start_time: float, ): - accumulator = AnthropicStreamAccumulator() + accumulator = _AnthropicStreamAccumulator() def generator(): try: diff --git a/posthog/ai/anthropic/anthropic_async.py b/posthog/ai/anthropic/anthropic_async.py index 3e2bf71cf..9ac091c64 100644 --- a/posthog/ai/anthropic/anthropic_async.py +++ b/posthog/ai/anthropic/anthropic_async.py @@ -11,10 +11,24 @@ from typing import Any, Dict, List, Optional from posthog import setup -from ..stream import AsyncStreamWrapper -from ..types import StreamingContentBlock, TokenUsage -from ..utils import call_llm_and_track_usage_async -from .anthropic_stream import AnthropicStreamAccumulator +from ..stream import AsyncStreamWrapper as AsyncStreamWrapper +from ..types import ( + StreamingContentBlock as StreamingContentBlock, + TokenUsage as TokenUsage, + ToolInProgress as ToolInProgress, +) +from ..utils import ( + call_llm_and_track_usage_async as call_llm_and_track_usage_async, + merge_usage_stats as merge_usage_stats, +) +from ._anthropic_stream import _AnthropicStreamAccumulator +from .anthropic_converter import ( + extract_anthropic_usage_from_event as extract_anthropic_usage_from_event, + finalize_anthropic_tool_input as finalize_anthropic_tool_input, + handle_anthropic_content_block_start as handle_anthropic_content_block_start, + handle_anthropic_text_delta as handle_anthropic_text_delta, + handle_anthropic_tool_delta as handle_anthropic_tool_delta, +) from posthog.client import Client as PostHogClient @@ -171,7 +185,7 @@ def _track_streaming_response( kwargs: Dict[str, Any], start_time: float, ): - accumulator = AnthropicStreamAccumulator() + accumulator = _AnthropicStreamAccumulator() async def generator(): try: