Skip to content

refactor(anthropic): share streaming state accumulation #885

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
marandaneto merged 2 commits into from
Aug 25, 2026
Merged

refactor(anthropic): share streaming state accumulation #885

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2928f7d
refactor(anthropic): share streaming accumulator
marandaneto Aug 21, 2026
7aca336
fix(anthropic): preserve public API after stream refactor
marandaneto Aug 22, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
60 changes: 60 additions & 0 deletions posthog/ai/anthropic/_anthropic_stream.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
96 changes: 21 additions & 75 deletions posthog/ai/anthropic/anthropic.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,17 +11,22 @@
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 ..types import (
StreamingContentBlock as StreamingContentBlock,
TokenUsage as TokenUsage,
ToolInProgress as ToolInProgress,
)
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 ..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
Expand Down Expand Up @@ -180,72 +185,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
Expand All @@ -257,11 +203,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)
Expand Down
98 changes: 22 additions & 76 deletions posthog/ai/anthropic/anthropic_async.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,18 +11,23 @@
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 ..stream import AsyncStreamWrapper as AsyncStreamWrapper
from ..types import (
StreamingContentBlock as StreamingContentBlock,
TokenUsage as TokenUsage,
ToolInProgress as ToolInProgress,
)
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 ..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

Expand Down Expand Up @@ -180,72 +185,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
Expand All @@ -257,11 +203,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)
Expand Down
Loading