"""PostHog LLM Analytics processor for the Claude Agent SDK. Wraps claude_agent_sdk.query() to automatically emit $ai_generation, $ai_span, and $ai_trace events to PostHog. """ import logging import time import uuid from dataclasses import dataclass, field from typing import Any, Callable, Dict, List, Optional, Union try: from claude_agent_sdk import ( AssistantMessage, ResultMessage, ToolUseBlock, UserMessage, ) from claude_agent_sdk import query as original_query from claude_agent_sdk.types import ClaudeAgentOptions, StreamEvent except ImportError: raise ModuleNotFoundError( "Please install the Claude Agent SDK to use this feature: 'pip install claude-agent-sdk'" ) from posthog import setup from posthog.client import Client log = logging.getLogger("posthog") @dataclass class _GenerationData: """Data accumulated for a single LLM generation (one API call).""" model: Optional[str] = None input_tokens: int = 0 output_tokens: int = 0 cache_read_input_tokens: int = 0 cache_creation_input_tokens: int = 0 start_time: float = 0.0 end_time: float = 0.0 span_id: str = field(default_factory=lambda: str(uuid.uuid4())) stop_reason: Optional[str] = None class _GenerationTracker: """Tracks StreamEvent boundaries to reconstruct per-generation metrics. Each message_start -> message_stop cycle in the Anthropic streaming protocol represents one API call (one generation). """ def __init__(self) -> None: self._current: Optional[_GenerationData] = None self._completed: List[_GenerationData] = [] self._last_model: Optional[str] = None self._received_stream_events: bool = False def process_stream_event(self, event: "StreamEvent") -> None: self._received_stream_events = True raw = event.event event_type = raw.get("type") if event_type == "message_start": self._current = _GenerationData(start_time=time.time()) message = raw.get("message", {}) self._current.model = message.get("model") usage = message.get("usage", {}) self._current.input_tokens = usage.get("input_tokens", 0) self._current.output_tokens = usage.get("output_tokens", 0) self._current.cache_read_input_tokens = usage.get( "cache_read_input_tokens", 0 ) self._current.cache_creation_input_tokens = usage.get( "cache_creation_input_tokens", 0 ) elif event_type == "message_delta" and self._current is not None: usage = raw.get("usage", {}) # message_delta usage reports cumulative output tokens if usage.get("output_tokens"): self._current.output_tokens = usage["output_tokens"] # Extract stop reason from message_delta delta_stop_reason = raw.get("delta", {}).get("stop_reason") if delta_stop_reason is not None: self._current.stop_reason = delta_stop_reason elif event_type == "message_stop" and self._current is not None: self._current.end_time = time.time() self._completed.append(self._current) self._last_model = self._current.model self._current = None def set_model(self, model: str) -> None: self._last_model = model @property def last_model(self) -> Optional[str]: return self._last_model def has_completed_generation(self) -> bool: return len(self._completed) > 0 def pop_generation(self) -> _GenerationData: return self._completed.pop(0) def has_pending(self) -> bool: return self._current is not None @property def generation_count(self) -> int: return len(self._completed) @property def current_span_id(self) -> Optional[str]: """Span ID of the generation currently in progress (before message_stop).""" return self._current.span_id if self._current else None @property def had_any_stream_events(self) -> bool: """Whether we received any StreamEvents at all.""" return self._received_stream_events class PostHogClaudeAgentProcessor: """Wraps claude_agent_sdk.query() to emit PostHog LLM analytics events. Emits: - $ai_generation: one per Anthropic API call (reconstructed from StreamEvents) - $ai_span: one per tool use (ToolUseBlock in AssistantMessage) - $ai_trace: one per query() call (on ResultMessage) """ def __init__( self, client: Optional[Client] = None, distinct_id: Optional[ Union[str, Callable[["ResultMessage"], Optional[str]]] ] = None, privacy_mode: bool = False, groups: Optional[Dict[str, Any]] = None, properties: Optional[Dict[str, Any]] = None, ): """ Initialize a Claude Agent SDK query processor. Args: client: Optional PostHog client. Uses the default client when omitted. distinct_id: Optional distinct ID for emitted events, or a callable that receives a ``ResultMessage`` and returns one. privacy_mode: Whether to redact captured inputs, outputs, and tool data. groups: Optional PostHog groups to associate with emitted events. properties: Additional properties included on emitted AI events. """ self._client = client or setup() self._distinct_id = distinct_id self._privacy_mode = privacy_mode self._groups = groups or {} self._properties = properties or {} def _get_distinct_id( self, result: Optional["ResultMessage"] = None ) -> Optional[str]: if callable(self._distinct_id): if result: val = self._distinct_id(result) if val: return str(val) return None elif self._distinct_id: return str(self._distinct_id) return None def _with_privacy_mode(self, value: Any) -> Any: if self._privacy_mode or ( hasattr(self._client, "privacy_mode") and self._client.privacy_mode ): return None return value def _capture_event( self, event: str, properties: Dict[str, Any], distinct_id: Optional[str] = None, groups: Optional[Dict[str, Any]] = None, ) -> None: try: if not hasattr(self._client, "capture") or not callable( self._client.capture ): return final_properties = { **properties, **self._properties, } self._client.capture( distinct_id=distinct_id or "unknown", event=event, properties=final_properties, groups=groups if groups is not None else self._groups, ) except Exception as e: log.debug(f"Failed to capture PostHog event: {e}") async def query( self, *, prompt: Any, options: Optional[ClaudeAgentOptions] = None, transport: Any = None, posthog_distinct_id: Optional[ Union[str, Callable[["ResultMessage"], Optional[str]]] ] = None, posthog_trace_id: Optional[str] = None, posthog_properties: Optional[Dict[str, Any]] = None, posthog_privacy_mode: Optional[bool] = None, posthog_groups: Optional[Dict[str, Any]] = None, ): """Drop-in replacement for claude_agent_sdk.query() with PostHog instrumentation. All original messages are yielded unchanged. PostHog events are emitted automatically in the background. Args: prompt: The prompt (same as claude_agent_sdk.query) options: ClaudeAgentOptions (same as claude_agent_sdk.query) transport: Optional transport (same as claude_agent_sdk.query) posthog_distinct_id: Override distinct_id for this query posthog_trace_id: Override trace_id for this query posthog_properties: Extra properties merged into all events for this query posthog_privacy_mode: Override privacy mode for this query posthog_groups: Override groups for this query """ from dataclasses import replace # Per-call overrides distinct_id_override = posthog_distinct_id or self._distinct_id trace_id = posthog_trace_id or str(uuid.uuid4()) extra_props = posthog_properties or {} privacy = ( posthog_privacy_mode if posthog_privacy_mode is not None else self._privacy_mode ) groups = posthog_groups or self._groups # Ensure partial messages are enabled for per-generation tracking if options is None: options = ClaudeAgentOptions(include_partial_messages=True) elif not options.include_partial_messages: options = replace(options, include_partial_messages=True) tracker = _GenerationTracker() query_start = time.time() generation_index = 0 current_generation_span_id: Optional[str] = None # Track input/output for generation events initial_input: List[Dict[str, Any]] = [] if isinstance(prompt, str): initial_input = [{"role": "user", "content": prompt}] if options and options.system_prompt and isinstance(options.system_prompt, str): initial_input = [ {"role": "system", "content": options.system_prompt} ] + initial_input # Two-slot input tracking: # - current_input: input for the generation currently in progress # - next_input: tool results that arrive mid-turn, queued for the next generation # # Message ordering from the SDK is: # message_start → content_blocks → AssistantMessage → UserMessage(tool result) → message_stop # So UserMessage arrives *before* message_stop. When message_stop fires we emit # with current_input, then promote next_input → current_input for the next turn. current_input: Optional[List[Dict[str, Any]]] = initial_input or None next_input: Optional[List[Dict[str, Any]]] = None # Accumulate assistant output per generation pending_output: List[Dict[str, Any]] = [] async for message in original_query( prompt=prompt, options=options, transport=transport ): # All instrumentation is wrapped in try/except so PostHog errors # never interrupt the underlying Claude Agent SDK query. try: if isinstance(message, StreamEvent): tracker.process_stream_event(message) # Emit $ai_generation when a turn completes if tracker.has_completed_generation(): gen = tracker.pop_generation() generation_index += 1 current_generation_span_id = gen.span_id self._emit_generation( gen, trace_id, generation_index, current_input, pending_output or None, distinct_id_override, extra_props, privacy, groups, ) # Promote: tool results from this turn become input for next turn current_input = next_input next_input = None pending_output = [] elif isinstance(message, AssistantMessage): tracker.set_model(message.model) # Use the in-progress generation's span_id as parent for tool spans. # AssistantMessage arrives before message_stop, so current_generation_span_id # would be stale (from the previous turn). tracker.current_span_id gives us # the correct in-progress generation. parent_id = tracker.current_span_id or current_generation_span_id # Build output content from assistant blocks output_content: List[Dict[str, Any]] = [] for block in message.content: if isinstance(block, ToolUseBlock): self._emit_tool_span( block, trace_id, parent_id, distinct_id_override, extra_props, privacy, groups, ) output_content.append( { "type": "function", "function": { "name": block.name, "arguments": block.input, }, } ) elif hasattr(block, "text"): output_content.append({"type": "text", "text": block.text}) if output_content: pending_output = [ {"role": "assistant", "content": output_content} ] elif isinstance(message, UserMessage): # UserMessages carry tool results. They arrive *before* message_stop # for the current turn, so queue them as input for the *next* generation. content = message.content if isinstance(content, str): next_input = [{"role": "user", "content": content}] elif isinstance(content, list): formatted: List[Dict[str, Any]] = [] for block in content: if hasattr(block, "tool_use_id"): formatted.append( { "type": "tool_result", "tool_use_id": block.tool_use_id, "content": str(block.content)[:500] if block.content else None, } ) elif hasattr(block, "text"): formatted.append({"type": "text", "text": block.text}) if formatted: next_input = [{"role": "user", "content": formatted}] elif isinstance(message, ResultMessage): # Fallback: if no StreamEvents were received, emit a single # generation from ResultMessage aggregate data if not tracker.had_any_stream_events: self._emit_generation_from_result( message, trace_id, tracker.last_model, query_start, initial_input, pending_output, distinct_id_override, extra_props, privacy, groups, ) self._emit_trace( message, trace_id, query_start, distinct_id_override, extra_props, privacy, groups, ) except Exception as e: log.debug(f"PostHog instrumentation error (non-fatal): {e}") yield message def _emit_generation( self, gen: _GenerationData, trace_id: str, generation_index: int, input_messages: Optional[List[Dict[str, Any]]], output_choices: Optional[List[Dict[str, Any]]], distinct_id: Any, extra_props: Dict[str, Any], privacy: bool, groups: Dict[str, Any], ) -> None: resolved_id = self._resolve_distinct_id(distinct_id) latency = ( (gen.end_time - gen.start_time) if gen.start_time and gen.end_time else 0 ) properties: Dict[str, Any] = { "$ai_trace_id": trace_id, "$ai_span_id": gen.span_id, "$ai_span_name": f"generation_{generation_index}", "$ai_provider": "anthropic", "$ai_framework": "claude-agent-sdk", "$ai_model": gen.model, "$ai_input_tokens": gen.input_tokens, "$ai_output_tokens": gen.output_tokens, "$ai_latency": latency, **extra_props, } if input_messages is not None: properties["$ai_input"] = ( None if privacy else self._with_privacy_mode(input_messages) ) if output_choices is not None: properties["$ai_output_choices"] = ( None if privacy else self._with_privacy_mode(output_choices) ) if gen.cache_read_input_tokens: properties["$ai_cache_read_input_tokens"] = gen.cache_read_input_tokens if gen.cache_creation_input_tokens: properties["$ai_cache_creation_input_tokens"] = ( gen.cache_creation_input_tokens ) if gen.stop_reason is not None: properties["$ai_stop_reason"] = gen.stop_reason if resolved_id is None: properties["$process_person_profile"] = False self._capture_event( "$ai_generation", properties, resolved_id or trace_id, groups ) def _emit_generation_from_result( self, result: "ResultMessage", trace_id: str, model: Optional[str], query_start: float, input_messages: Optional[List[Dict[str, Any]]], output_choices: Optional[List[Dict[str, Any]]], distinct_id: Any, extra_props: Dict[str, Any], privacy: bool, groups: Dict[str, Any], ) -> None: """Fallback: emit a single generation from ResultMessage aggregate data.""" resolved_id = self._resolve_distinct_id(distinct_id) usage = result.usage or {} properties: Dict[str, Any] = { "$ai_trace_id": trace_id, "$ai_span_id": str(uuid.uuid4()), "$ai_span_name": "generation_1", "$ai_provider": "anthropic", "$ai_framework": "claude-agent-sdk", "$ai_model": model, "$ai_input_tokens": usage.get("input_tokens", 0), "$ai_output_tokens": usage.get("output_tokens", 0), "$ai_latency": result.duration_api_ms / 1000.0 if result.duration_api_ms else 0, "$ai_is_error": result.is_error, **extra_props, } if input_messages is not None: properties["$ai_input"] = ( None if privacy else self._with_privacy_mode(input_messages) ) if output_choices is not None: properties["$ai_output_choices"] = ( None if privacy else self._with_privacy_mode(output_choices) ) cache_read = usage.get("cache_read_input_tokens", 0) cache_creation = usage.get("cache_creation_input_tokens", 0) if cache_read: properties["$ai_cache_read_input_tokens"] = cache_read if cache_creation: properties["$ai_cache_creation_input_tokens"] = cache_creation if result.total_cost_usd is not None: properties["$ai_total_cost_usd"] = result.total_cost_usd if resolved_id is None: properties["$process_person_profile"] = False self._capture_event( "$ai_generation", properties, resolved_id or trace_id, groups ) def _emit_tool_span( self, block: "ToolUseBlock", trace_id: str, parent_span_id: Optional[str], distinct_id: Any, extra_props: Dict[str, Any], privacy: bool, groups: Dict[str, Any], ) -> None: resolved_id = self._resolve_distinct_id(distinct_id) properties: Dict[str, Any] = { "$ai_trace_id": trace_id, "$ai_span_id": str(uuid.uuid4()), "$ai_parent_id": parent_span_id, "$ai_span_name": block.name, "$ai_span_type": "tool", "$ai_provider": "anthropic", "$ai_framework": "claude-agent-sdk", **extra_props, } if not privacy and not ( hasattr(self._client, "privacy_mode") and self._client.privacy_mode ): properties["$ai_input_state"] = _ensure_serializable(block.input) if resolved_id is None: properties["$process_person_profile"] = False self._capture_event("$ai_span", properties, resolved_id or trace_id, groups) def _emit_trace( self, result: "ResultMessage", trace_id: str, query_start: float, distinct_id: Any, extra_props: Dict[str, Any], privacy: bool, groups: Dict[str, Any], ) -> None: resolved_id = self._resolve_distinct_id(distinct_id, result) latency = ( result.duration_ms / 1000.0 if result.duration_ms else (time.time() - query_start) ) properties: Dict[str, Any] = { "$ai_trace_id": trace_id, "$ai_trace_name": "claude_agent_sdk_query", "$ai_provider": "anthropic", "$ai_framework": "claude-agent-sdk", "$ai_latency": latency, "$ai_is_error": result.is_error, **extra_props, } if result.total_cost_usd is not None: properties["$ai_total_cost_usd"] = result.total_cost_usd if resolved_id is None: properties["$process_person_profile"] = False self._capture_event("$ai_trace", properties, resolved_id or trace_id, groups) # Flush to ensure events are sent before process exits try: if hasattr(self._client, "flush") and callable(self._client.flush): self._client.flush() except Exception as e: log.debug(f"Error flushing PostHog client: {e}") def _resolve_distinct_id( self, override: Any, result: Optional["ResultMessage"] = None, ) -> Optional[str]: """Resolve distinct_id from override or instance default.""" if callable(override): if result: val = override(result) if val: return str(val) return None elif override: return str(override) # Fall back to instance default return self._get_distinct_id(result) def _ensure_serializable(obj: Any) -> Any: """Ensure an object is JSON-serializable.""" if obj is None: return None try: import json json.dumps(obj) return obj except (TypeError, ValueError): return str(obj)