From cd07bd09164edb564356c90d056dc426ca69f33f Mon Sep 17 00:00:00 2001 From: Andrew Maguire Date: Tue, 2 Dec 2025 16:16:08 +0000 Subject: [PATCH 01/14] feat: Add Pydantic AI instrumentation via OpenTelemetry Add OpenTelemetry-based instrumentation for Pydantic AI agents: - PostHogSpanExporter: Generic OTel exporter that translates spans to PostHog events - PydanticAISpanExporter: Wrapper that normalizes Pydantic AI message formats and tool attributes - instrument_pydantic_ai(): One-liner setup function for easy integration Usage: from posthog.ai.pydantic_ai import instrument_pydantic_ai instrument_pydantic_ai(posthog_client, distinct_id="user_123") agent = Agent("openai:gpt-4") result = await agent.run("Hello") # Automatically traced Co-Authored-By: Claude --- posthog/ai/otel/README.md | 84 +++ posthog/ai/otel/__init__.py | 10 + posthog/ai/otel/exporter.py | 532 +++++++++++++++++ posthog/ai/pydantic_ai/README.md | 134 +++++ posthog/ai/pydantic_ai/__init__.py | 10 + posthog/ai/pydantic_ai/exporter.py | 259 +++++++++ posthog/ai/pydantic_ai/instrument.py | 92 +++ posthog/test/ai/otel/__init__.py | 0 posthog/test/ai/otel/test_exporter.py | 440 ++++++++++++++ posthog/test/ai/pydantic_ai/__init__.py | 0 posthog/test/ai/pydantic_ai/test_exporter.py | 537 ++++++++++++++++++ .../test/ai/pydantic_ai/test_instrument.py | 97 ++++ pyproject.toml | 6 + 13 files changed, 2201 insertions(+) create mode 100644 posthog/ai/otel/README.md create mode 100644 posthog/ai/otel/__init__.py create mode 100644 posthog/ai/otel/exporter.py create mode 100644 posthog/ai/pydantic_ai/README.md create mode 100644 posthog/ai/pydantic_ai/__init__.py create mode 100644 posthog/ai/pydantic_ai/exporter.py create mode 100644 posthog/ai/pydantic_ai/instrument.py create mode 100644 posthog/test/ai/otel/__init__.py create mode 100644 posthog/test/ai/otel/test_exporter.py create mode 100644 posthog/test/ai/pydantic_ai/__init__.py create mode 100644 posthog/test/ai/pydantic_ai/test_exporter.py create mode 100644 posthog/test/ai/pydantic_ai/test_instrument.py diff --git a/posthog/ai/otel/README.md b/posthog/ai/otel/README.md new file mode 100644 index 00000000..aef22086 --- /dev/null +++ b/posthog/ai/otel/README.md @@ -0,0 +1,84 @@ +# PostHog OpenTelemetry Integration + +This module provides a generic OpenTelemetry `SpanExporter` that translates OTel spans into PostHog AI analytics events. + +## Overview + +Many AI/LLM frameworks use OpenTelemetry for instrumentation. This exporter allows PostHog to receive telemetry from any OTel-instrumented framework by converting spans to PostHog events. + +``` +┌─────────────────┐ ┌──────────────┐ ┌─────────────────────┐ ┌─────────┐ +│ AI Framework │────>│ OTel Spans │────>│ PostHogSpanExporter │────>│ PostHog │ +│ (Pydantic AI, │ │ (native) │ │ (translates spans) │ │ Events │ +│ LlamaIndex...) │ └──────────────┘ └─────────────────────┘ └─────────┘ +└─────────────────┘ +``` + +## Usage + +```python +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import BatchSpanProcessor +from posthog import Posthog +from posthog.ai.otel import PostHogSpanExporter + +# Create PostHog client +posthog = Posthog(project_api_key="phc_xxx", host="https://us.i.posthog.com") + +# Create exporter and tracer provider +exporter = PostHogSpanExporter( + client=posthog, + distinct_id="user_123", + privacy_mode=False, # Set True to exclude message content +) + +provider = TracerProvider() +provider.add_span_processor(BatchSpanProcessor(exporter)) + +# Use this provider with your OTel-instrumented framework +``` + +## Span to Event Mapping + +The exporter classifies spans and maps them to PostHog events: + +| Span Type | PostHog Event | Detection | +|-----------|---------------|-----------| +| Model request | `$ai_generation` | Span name starts with "chat" or has `gen_ai.request.model` | +| Tool execution | `$ai_span` | Span name contains "tool" or has `gen_ai.tool.name` | +| Agent orchestration | (skipped) | Span name contains "agent" | + +## GenAI Semantic Conventions + +The exporter follows [OpenTelemetry GenAI semantic conventions](https://opentelemetry.io/docs/specs/semconv/gen-ai/): + +| OTel Attribute | PostHog Property | +|----------------|------------------| +| `gen_ai.request.model` | `$ai_model` | +| `gen_ai.system` | `$ai_provider` | +| `gen_ai.usage.input_tokens` | `$ai_input_tokens` | +| `gen_ai.usage.output_tokens` | `$ai_output_tokens` | +| `gen_ai.input.messages` | `$ai_input` | +| `gen_ai.output.messages` | `$ai_output_choices` | +| `gen_ai.tool.name` | `$ai_span_name` | +| `gen_ai.tool.call.arguments` | `$ai_tool_arguments` | +| `gen_ai.tool.call.result` | `$ai_tool_result` | + +## Configuration Options + +| Parameter | Type | Description | +|-----------|------|-------------| +| `client` | `Posthog` | PostHog client instance | +| `distinct_id` | `str` | User identifier (falls back to trace ID if not set) | +| `privacy_mode` | `bool` | Exclude message content from events | +| `properties` | `dict` | Additional properties to include in all events | +| `groups` | `dict` | PostHog groups for all events | +| `debug` | `bool` | Enable debug logging | + +## Framework-Specific Exporters + +For frameworks with non-standard attribute names or message formats, use the framework-specific exporter wrapper: + +- **Pydantic AI**: Use `posthog.ai.pydantic_ai.PydanticAISpanExporter` or the simpler `instrument_pydantic_ai()` function + +These wrappers normalize framework-specific formats before passing spans to `PostHogSpanExporter`. diff --git a/posthog/ai/otel/__init__.py b/posthog/ai/otel/__init__.py new file mode 100644 index 00000000..ed6d52b6 --- /dev/null +++ b/posthog/ai/otel/__init__.py @@ -0,0 +1,10 @@ +""" +OpenTelemetry integration for PostHog AI observability. + +This module provides a SpanExporter that translates OpenTelemetry spans +(particularly GenAI semantic convention spans) into PostHog AI events. +""" + +from posthog.ai.otel.exporter import PostHogSpanExporter + +__all__ = ["PostHogSpanExporter"] diff --git a/posthog/ai/otel/exporter.py b/posthog/ai/otel/exporter.py new file mode 100644 index 00000000..de12666a --- /dev/null +++ b/posthog/ai/otel/exporter.py @@ -0,0 +1,532 @@ +""" +PostHog SpanExporter for OpenTelemetry. + +Translates OpenTelemetry spans (using GenAI semantic conventions) into PostHog AI events. +This enables any OTel-instrumented AI framework (Pydantic AI, LlamaIndex, etc.) to send +telemetry to PostHog. +""" + +import json +import logging +from typing import Any, Dict, Optional, Sequence, Union + +from posthog.client import Client as PostHogClient + +try: + from opentelemetry.sdk.trace import ReadableSpan + from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult + from opentelemetry.trace import StatusCode + + OTEL_AVAILABLE = True +except ImportError: + OTEL_AVAILABLE = False + # Define stub types for type hints when OTel is not installed + ReadableSpan = Any # type: ignore + SpanExporter = object # type: ignore + SpanExportResult = Any # type: ignore + StatusCode = Any # type: ignore + +logger = logging.getLogger(__name__) + + +# OTel GenAI semantic convention attribute names +# See: https://opentelemetry.io/docs/specs/semconv/gen-ai/ +class GenAIAttributes: + # Operation + OPERATION_NAME = "gen_ai.operation.name" + + # Request attributes + REQUEST_MODEL = "gen_ai.request.model" + SYSTEM = "gen_ai.system" + PROVIDER_NAME = "gen_ai.provider_name" # Alternative to gen_ai.system + + # Response attributes + RESPONSE_MODEL = "gen_ai.response.model" + RESPONSE_ID = "gen_ai.response.id" + FINISH_REASONS = "gen_ai.response.finish_reasons" + + # Usage attributes + INPUT_TOKENS = "gen_ai.usage.input_tokens" + OUTPUT_TOKENS = "gen_ai.usage.output_tokens" + + # Message content (when captured) + INPUT_MESSAGES = "gen_ai.input.messages" + OUTPUT_MESSAGES = "gen_ai.output.messages" + SYSTEM_INSTRUCTIONS = "gen_ai.system_instructions" + + # Pydantic AI specific + AGENT_NAME = "gen_ai.agent.name" + AGENT_NAME_LEGACY = "agent_name" + + # Tool attributes + TOOL_NAME = "gen_ai.tool.name" + TOOL_CALL_ID = "gen_ai.tool.call.id" + TOOL_ARGUMENTS = "gen_ai.tool.call.arguments" + TOOL_RESULT = "gen_ai.tool.call.result" + + # Model parameters + TEMPERATURE = "gen_ai.request.temperature" + TOP_P = "gen_ai.request.top_p" + MAX_TOKENS = "gen_ai.request.max_tokens" + FREQUENCY_PENALTY = "gen_ai.request.frequency_penalty" + PRESENCE_PENALTY = "gen_ai.request.presence_penalty" + SEED = "gen_ai.request.seed" + + # Server info + SERVER_ADDRESS = "server.address" + SERVER_PORT = "server.port" + + +class PostHogSpanExporter(SpanExporter if OTEL_AVAILABLE else object): + """ + OpenTelemetry SpanExporter that sends AI/LLM spans to PostHog. + + Translates OTel GenAI semantic convention spans into PostHog AI events: + - Model request spans → $ai_generation + - Agent run spans → $ai_trace + - Tool execution spans → $ai_span + + Usage: + from posthog import Posthog + from posthog.ai.otel import PostHogSpanExporter + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import BatchSpanProcessor + + posthog = Posthog(api_key="...", host="...") + exporter = PostHogSpanExporter(posthog, distinct_id="user_123") + + provider = TracerProvider() + provider.add_span_processor(BatchSpanProcessor(exporter)) + """ + + def __init__( + self, + client: PostHogClient, + distinct_id: Optional[str] = None, + privacy_mode: bool = False, + properties: Optional[Dict[str, Any]] = None, + groups: Optional[Dict[str, Any]] = None, + debug: bool = False, + ): + """ + Initialize the PostHog span exporter. + + Args: + client: PostHog client instance + distinct_id: Default distinct ID for events (can be overridden per-span) + privacy_mode: If True, redact message content from events + properties: Additional properties to include in all events + groups: PostHog groups for all events + debug: Enable debug logging + """ + if not OTEL_AVAILABLE: + raise ImportError( + "OpenTelemetry SDK is required for PostHogSpanExporter. " + "Install it with: pip install opentelemetry-sdk" + ) + + self._client = client + self._distinct_id = distinct_id + self._privacy_mode = privacy_mode or getattr(client, "privacy_mode", False) + self._properties = properties or {} + self._groups = groups + self._debug = debug + + def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: + """ + Export spans to PostHog. + + Translates each span into the appropriate PostHog event type. + """ + for span in spans: + try: + event = self._span_to_event(span) + if event: + distinct_id = self._get_distinct_id(span) + + if self._debug: + logger.debug( + f"Exporting span '{span.name}' as {event['name']} " + f"with distinct_id={distinct_id}" + ) + + capture_kwargs: Dict[str, Any] = { + "distinct_id": distinct_id, + "event": event["name"], + "properties": event["properties"], + } + + if self._groups: + capture_kwargs["groups"] = self._groups + + self._client.capture(**capture_kwargs) + + except Exception as e: + logger.warning(f"Failed to export span '{span.name}': {e}") + if self._debug: + logger.exception("Full exception:") + + return SpanExportResult.SUCCESS + + def shutdown(self) -> None: + """Shutdown the exporter.""" + pass + + def force_flush(self, timeout_millis: int = 30000) -> bool: + """Force flush any buffered spans.""" + return True + + def _get_distinct_id(self, span: ReadableSpan) -> str: + """Get distinct ID for a span, with fallback to trace ID.""" + attrs = dict(span.attributes or {}) + + # Check for custom distinct_id attribute + distinct_id = attrs.get("posthog.distinct_id") + if distinct_id: + return str(distinct_id) + + # Use configured default + if self._distinct_id: + return self._distinct_id + + # Fall back to trace ID + return format(span.context.trace_id, "032x") + + def _span_to_event(self, span: ReadableSpan) -> Optional[Dict[str, Any]]: + """ + Convert an OTel span to a PostHog event. + + Returns None for spans that shouldn't be exported. + """ + attrs = dict(span.attributes or {}) + span_name = span.name + + # Calculate latency in seconds + latency = ( + (span.end_time - span.start_time) / 1e9 + if span.end_time and span.start_time + else 0 + ) + + # Format trace ID as UUID (with dashes) for PostHog compatibility + trace_id = self._format_trace_id_as_uuid(span.context.trace_id) + # Span IDs remain as hex (no dashes needed) + span_id = format(span.context.span_id, "016x") + parent_span_id = ( + format(span.parent.span_id, "016x") if span.parent else None + ) + + # Check for error status + is_error = span.status.status_code == StatusCode.ERROR if span.status else False + error_message = span.status.description if is_error and span.status else None + + # Model request span → $ai_generation + if self._is_generation_span(span_name, attrs): + return self._create_generation_event( + span, attrs, trace_id, span_id, parent_span_id, latency, is_error, error_message + ) + + # Agent run span → skip (PostHog UI auto-creates trace wrapper from generation events) + # The $ai_trace_id on generation events is sufficient for grouping + if self._is_agent_span(span_name, attrs): + return None # Don't emit separate $ai_trace events + + # Tool execution span → $ai_span + if self._is_tool_span(span_name, attrs): + return self._create_tool_span_event( + span, attrs, trace_id, span_id, parent_span_id, latency, is_error, error_message + ) + + # Generic span that might be part of AI workflow + if self._is_ai_related_span(span_name, attrs): + return self._create_span_event( + span, attrs, trace_id, span_id, parent_span_id, latency, is_error, error_message + ) + + return None + + def _is_generation_span(self, span_name: str, attrs: Dict[str, Any]) -> bool: + """Check if span represents an LLM generation/chat completion.""" + operation = attrs.get(GenAIAttributes.OPERATION_NAME, "") + return ( + span_name.startswith("chat ") + or operation == "chat" + or attrs.get(GenAIAttributes.REQUEST_MODEL) is not None + ) + + def _is_agent_span(self, span_name: str, attrs: Dict[str, Any]) -> bool: + """Check if span represents an agent run.""" + return span_name in ("agent run", "invoke_agent") or attrs.get( + GenAIAttributes.AGENT_NAME + ) + + def _is_tool_span(self, span_name: str, attrs: Dict[str, Any]) -> bool: + """Check if span represents a tool/function execution.""" + return ( + "tool" in span_name.lower() + or "execute_tool" in span_name + or attrs.get(GenAIAttributes.TOOL_NAME) is not None + ) + + def _is_ai_related_span(self, span_name: str, attrs: Dict[str, Any]) -> bool: + """Check if span is AI-related based on attributes.""" + ai_attrs = [ + GenAIAttributes.SYSTEM, + GenAIAttributes.PROVIDER_NAME, + GenAIAttributes.REQUEST_MODEL, + GenAIAttributes.AGENT_NAME, + ] + return any(attrs.get(attr) for attr in ai_attrs) + + def _get_generation_span_name(self, span_name: str, provider: str) -> str: + """ + Derive a descriptive span name for generation events. + + Returns something like 'openai_chat_completions' based on provider. + """ + # If span name already looks like a good identifier, use it + if span_name and not span_name.startswith("chat "): + # Clean up span name to be a good identifier + clean_name = span_name.replace(" ", "_").replace("-", "_").lower() + return clean_name + + # Otherwise derive from provider + provider_lower = str(provider).lower() if provider else "unknown" + return f"{provider_lower}_chat_completions" + + def _create_generation_event( + self, + span: ReadableSpan, + attrs: Dict[str, Any], + trace_id: str, + span_id: str, + parent_span_id: Optional[str], + latency: float, + is_error: bool, + error_message: Optional[str], + ) -> Dict[str, Any]: + """Create a $ai_generation event from a model request span.""" + # Extract model and provider info + model = attrs.get(GenAIAttributes.REQUEST_MODEL) or attrs.get( + GenAIAttributes.RESPONSE_MODEL + ) + provider = attrs.get(GenAIAttributes.SYSTEM) or attrs.get( + GenAIAttributes.PROVIDER_NAME, "unknown" + ) + + # Extract token usage + input_tokens = attrs.get(GenAIAttributes.INPUT_TOKENS) + output_tokens = attrs.get(GenAIAttributes.OUTPUT_TOKENS) + + # Extract messages (respecting privacy mode) + input_messages = None + output_messages = None + if not self._privacy_mode: + input_messages = self._parse_json_attr( + attrs.get(GenAIAttributes.INPUT_MESSAGES) + ) + output_messages = self._parse_json_attr( + attrs.get(GenAIAttributes.OUTPUT_MESSAGES) + ) + + # Build base URL from server info + server_address = attrs.get(GenAIAttributes.SERVER_ADDRESS) + server_port = attrs.get(GenAIAttributes.SERVER_PORT) + base_url = None + if server_address: + base_url = f"https://{server_address}" + if server_port: + base_url = f"{base_url}:{server_port}" + + # Extract model parameters + model_params = {} + param_attrs = [ + (GenAIAttributes.TEMPERATURE, "temperature"), + (GenAIAttributes.TOP_P, "top_p"), + (GenAIAttributes.MAX_TOKENS, "max_tokens"), + (GenAIAttributes.FREQUENCY_PENALTY, "frequency_penalty"), + (GenAIAttributes.PRESENCE_PENALTY, "presence_penalty"), + (GenAIAttributes.SEED, "seed"), + ] + for otel_attr, param_name in param_attrs: + if otel_attr in attrs: + model_params[param_name] = attrs[otel_attr] + + # Derive span name from span name or provider + generation_span_name = self._get_generation_span_name(span.name, provider) + + # PostHog expects generation events to NOT have span_id/parent_id + # The $ai_trace_id alone is sufficient for grouping + properties: Dict[str, Any] = { + "$ai_trace_id": trace_id, + "$ai_span_name": generation_span_name, + "$ai_model": model, + "$ai_provider": provider, + "$ai_latency": latency, + "$ai_http_status": 500 if is_error else 200, + "$ai_is_error": is_error, + "$ai_framework": "pydantic-ai", + **self._properties, + } + + if model_params: + properties["$ai_model_parameters"] = model_params + + if input_tokens is not None: + properties["$ai_input_tokens"] = input_tokens + + if output_tokens is not None: + properties["$ai_output_tokens"] = output_tokens + + if input_messages is not None: + properties["$ai_input"] = input_messages + + if output_messages is not None: + properties["$ai_output_choices"] = output_messages + + if base_url: + properties["$ai_base_url"] = base_url + + if is_error and error_message: + properties["$ai_error"] = error_message + + # Handle distinct_id for person profile processing + if not self._distinct_id and not attrs.get("posthog.distinct_id"): + properties["$process_person_profile"] = False + + return {"name": "$ai_generation", "properties": properties} + + def _create_trace_event( + self, + span: ReadableSpan, + attrs: Dict[str, Any], + trace_id: str, + span_id: str, + latency: float, + is_error: bool, + error_message: Optional[str], + ) -> Dict[str, Any]: + """Create a $ai_trace event from an agent run span.""" + agent_name = attrs.get(GenAIAttributes.AGENT_NAME) or attrs.get( + GenAIAttributes.AGENT_NAME_LEGACY, "unknown" + ) + + properties: Dict[str, Any] = { + "$ai_trace_id": trace_id, + "$ai_span_id": span_id, + "$ai_span_name": agent_name, + "$ai_latency": latency, + "$ai_is_error": is_error, + "$ai_framework": "pydantic-ai", + **self._properties, + } + + if is_error and error_message: + properties["$ai_error"] = error_message + + if not self._distinct_id and not attrs.get("posthog.distinct_id"): + properties["$process_person_profile"] = False + + return {"name": "$ai_trace", "properties": properties} + + def _create_tool_span_event( + self, + span: ReadableSpan, + attrs: Dict[str, Any], + trace_id: str, + span_id: str, + parent_span_id: Optional[str], + latency: float, + is_error: bool, + error_message: Optional[str], + ) -> Dict[str, Any]: + """Create a $ai_span event from a tool execution span.""" + tool_name = attrs.get(GenAIAttributes.TOOL_NAME, span.name) + + properties: Dict[str, Any] = { + "$ai_trace_id": trace_id, + "$ai_span_id": span_id, + "$ai_span_name": tool_name, + "$ai_latency": latency, + "$ai_is_error": is_error, + "$ai_framework": "pydantic-ai", + **self._properties, + } + + if parent_span_id: + properties["$ai_parent_id"] = parent_span_id + + # Include tool arguments and result if not in privacy mode + if not self._privacy_mode: + tool_args = attrs.get(GenAIAttributes.TOOL_ARGUMENTS) + if tool_args: + properties["$ai_tool_arguments"] = self._parse_json_attr(tool_args) + + tool_result = attrs.get(GenAIAttributes.TOOL_RESULT) + if tool_result: + properties["$ai_tool_result"] = self._parse_json_attr(tool_result) + + if is_error and error_message: + properties["$ai_error"] = error_message + + if not self._distinct_id and not attrs.get("posthog.distinct_id"): + properties["$process_person_profile"] = False + + return {"name": "$ai_span", "properties": properties} + + def _create_span_event( + self, + span: ReadableSpan, + attrs: Dict[str, Any], + trace_id: str, + span_id: str, + parent_span_id: Optional[str], + latency: float, + is_error: bool, + error_message: Optional[str], + ) -> Dict[str, Any]: + """Create a generic $ai_span event.""" + properties: Dict[str, Any] = { + "$ai_trace_id": trace_id, + "$ai_span_id": span_id, + "$ai_span_name": span.name, + "$ai_latency": latency, + "$ai_is_error": is_error, + "$ai_framework": "pydantic-ai", + **self._properties, + } + + if parent_span_id: + properties["$ai_parent_id"] = parent_span_id + + if is_error and error_message: + properties["$ai_error"] = error_message + + if not self._distinct_id and not attrs.get("posthog.distinct_id"): + properties["$process_person_profile"] = False + + return {"name": "$ai_span", "properties": properties} + + def _parse_json_attr( + self, value: Optional[Union[str, Any]] + ) -> Optional[Any]: + """Parse a JSON string attribute, returning the value as-is if already parsed.""" + if value is None: + return None + if isinstance(value, str): + try: + return json.loads(value) + except json.JSONDecodeError: + return value + return value + + def _format_trace_id_as_uuid(self, trace_id: int) -> str: + """ + Convert an OTel trace ID (128-bit int) to UUID format with dashes. + + PostHog expects trace IDs in UUID format (e.g., 'a8f3d2c4-1247-4c40-8342-23d5e8d52584') + but OTel uses 128-bit integers formatted as 32 hex chars without dashes. + """ + hex_str = format(trace_id, "032x") + # Insert dashes to form UUID format: 8-4-4-4-12 + return f"{hex_str[:8]}-{hex_str[8:12]}-{hex_str[12:16]}-{hex_str[16:20]}-{hex_str[20:]}" diff --git a/posthog/ai/pydantic_ai/README.md b/posthog/ai/pydantic_ai/README.md new file mode 100644 index 00000000..edb3fc06 --- /dev/null +++ b/posthog/ai/pydantic_ai/README.md @@ -0,0 +1,134 @@ +# PostHog Pydantic AI Integration + +This module provides PostHog instrumentation for [Pydantic AI](https://ai.pydantic.dev/) agents. + +## Quick Start + +```python +from posthog import Posthog +from posthog.ai.pydantic_ai import instrument_pydantic_ai +from pydantic_ai import Agent + +# Initialize PostHog +posthog = Posthog(project_api_key="phc_xxx", host="https://us.i.posthog.com") + +# Instrument all Pydantic AI agents (call once at startup) +instrument_pydantic_ai(posthog, distinct_id="user_123") + +# Use Pydantic AI normally - all calls are automatically traced +agent = Agent("openai:gpt-4o") +result = await agent.run("What's the weather in San Francisco?") +``` + +## How It Works + +``` +┌─────────────────┐ ┌──────────────┐ ┌────────────────────────┐ ┌─────────┐ +│ agent.run() │────>│ OTel Spans │────>│ PydanticAISpanExporter │────>│ PostHog │ +│ │ │ (Pydantic AI │ │ (normalizes messages, │ │ Events │ +│ │ │ native) │ │ maps tool attributes) │ │ │ +└─────────────────┘ └──────────────┘ └────────────────────────┘ └─────────┘ +``` + +1. **Pydantic AI** emits OpenTelemetry spans natively via `Agent.instrument_all()` +2. **PydanticAISpanExporter** transforms Pydantic-specific formats to standard GenAI conventions +3. **PostHogSpanExporter** converts spans to PostHog `$ai_generation` and `$ai_span` events + +## Configuration Options + +```python +instrument_pydantic_ai( + client=posthog, # PostHog client instance + distinct_id="user_123", # User identifier for events + privacy_mode=False, # Set True to exclude message content + properties={ # Additional properties for all events + "$ai_session_id": "session_abc", + }, + groups={ # PostHog groups + "company": "acme", + }, + debug=False, # Enable debug logging +) +``` + +## What Gets Captured + +### Model Calls (`$ai_generation` events) + +Every LLM API call creates an event with: +- Model name and provider +- Input/output messages (unless `privacy_mode=True`) +- Token usage (input, output) +- Latency +- Error status + +### Tool Calls (`$ai_span` events) + +When agents use tools: +- Tool name +- Arguments passed to the tool +- Tool result/response +- Latency + +## Pydantic AI-Specific Handling + +This integration handles Pydantic AI's specific message and attribute formats: + +### Message Normalization + +Pydantic AI uses a "parts" format for messages: +```python +# Pydantic AI format +{"parts": [{"content": "Hello", "type": "text"}], "role": "user"} + +# Normalized to OpenAI format for PostHog +{"content": "Hello", "role": "user"} +``` + +### Tool Attribute Mapping + +Pydantic AI uses non-standard attribute names: +```python +# Pydantic AI attributes +"tool_arguments": '{"city": "SF"}' +"tool_response": "Sunny, 72F" + +# Mapped to GenAI standard +"gen_ai.tool.call.arguments": '{"city": "SF"}' +"gen_ai.tool.call.result": "Sunny, 72F" +``` + +## Requirements + +- `pydantic-ai >= 0.1.0` +- `opentelemetry-sdk` + +Install with: +```bash +pip install posthog pydantic-ai opentelemetry-sdk +``` + +## Advanced: Using the Exporter Directly + +For more control, use `PydanticAISpanExporter` directly: + +```python +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import BatchSpanProcessor +from pydantic_ai import Agent +from pydantic_ai.models.instrumented import InstrumentationSettings +from posthog.ai.pydantic_ai import PydanticAISpanExporter + +exporter = PydanticAISpanExporter( + client=posthog, + distinct_id="user_123", +) + +provider = TracerProvider() +provider.add_span_processor(BatchSpanProcessor(exporter)) + +Agent.instrument_all(InstrumentationSettings( + tracer_provider=provider, + include_content=True, +)) +``` diff --git a/posthog/ai/pydantic_ai/__init__.py b/posthog/ai/pydantic_ai/__init__.py new file mode 100644 index 00000000..e670c503 --- /dev/null +++ b/posthog/ai/pydantic_ai/__init__.py @@ -0,0 +1,10 @@ +""" +Pydantic AI integration for PostHog AI observability. + +This module provides a simple interface to instrument Pydantic AI agents +with PostHog tracing. +""" + +from posthog.ai.pydantic_ai.instrument import instrument_pydantic_ai + +__all__ = ["instrument_pydantic_ai"] diff --git a/posthog/ai/pydantic_ai/exporter.py b/posthog/ai/pydantic_ai/exporter.py new file mode 100644 index 00000000..85095c11 --- /dev/null +++ b/posthog/ai/pydantic_ai/exporter.py @@ -0,0 +1,259 @@ +""" +Pydantic AI specific SpanExporter for PostHog. + +This exporter wraps the generic PostHogSpanExporter and handles +Pydantic AI-specific transformations like message format normalization. +""" + +from typing import Any, Dict, List, Optional, Sequence + +try: + from opentelemetry.sdk.trace import ReadableSpan + from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult + + OTEL_AVAILABLE = True +except ImportError: + OTEL_AVAILABLE = False + ReadableSpan = Any # type: ignore + SpanExporter = object # type: ignore + SpanExportResult = Any # type: ignore + +from posthog.ai.otel import PostHogSpanExporter +from posthog.client import Client as PostHogClient + + +class PydanticAISpanExporter(SpanExporter if OTEL_AVAILABLE else object): + """ + SpanExporter for Pydantic AI that normalizes messages to OpenAI format. + + Pydantic AI uses its own message format with "parts": + {"parts": [{"content": "...", "type": "text"}], "role": "user"} + + This exporter transforms that to the standard OpenAI format: + {"content": "...", "role": "user"} + + This ensures consistent display in PostHog's LLM Analytics UI. + """ + + def __init__( + self, + client: PostHogClient, + distinct_id: Optional[str] = None, + privacy_mode: bool = False, + properties: Optional[Dict[str, Any]] = None, + groups: Optional[Dict[str, Any]] = None, + debug: bool = False, + ): + if not OTEL_AVAILABLE: + raise ImportError( + "OpenTelemetry SDK is required. Install with: pip install opentelemetry-sdk" + ) + + # Wrap the generic PostHog exporter + self._base_exporter = PostHogSpanExporter( + client=client, + distinct_id=distinct_id, + privacy_mode=privacy_mode, + properties=properties, + groups=groups, + debug=debug, + ) + self._debug = debug + + def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: + """Export spans after normalizing Pydantic AI message formats.""" + # Transform spans to normalize message format + transformed_spans = [self._transform_span(span) for span in spans] + return self._base_exporter.export(transformed_spans) + + def shutdown(self) -> None: + """Shutdown the exporter.""" + self._base_exporter.shutdown() + + def force_flush(self, timeout_millis: int = 30000) -> bool: + """Force flush any buffered spans.""" + return self._base_exporter.force_flush(timeout_millis) + + def _transform_span(self, span: ReadableSpan) -> ReadableSpan: + """ + Transform a span's attributes to normalize Pydantic AI-specific formats. + + Handles: + - Message format: Pydantic AI "parts" format → OpenAI format + - Tool attributes: tool_arguments/tool_response → gen_ai.tool.call.arguments/result + """ + import json + + attrs = dict(span.attributes or {}) + modified = False + + # Normalize input messages from Pydantic AI "parts" format + input_msgs = attrs.get("gen_ai.input.messages") + if input_msgs: + normalized = self._normalize_messages(input_msgs) + if normalized != input_msgs: + attrs["gen_ai.input.messages"] = ( + json.dumps(normalized) if isinstance(normalized, list) else normalized + ) + modified = True + + # Normalize output messages from Pydantic AI "parts" format + output_msgs = attrs.get("gen_ai.output.messages") + if output_msgs: + normalized = self._normalize_messages(output_msgs) + if normalized != output_msgs: + attrs["gen_ai.output.messages"] = ( + json.dumps(normalized) if isinstance(normalized, list) else normalized + ) + modified = True + + # Map Pydantic AI tool attributes to GenAI standard names + # Pydantic AI uses: tool_arguments, tool_response + # GenAI standard: gen_ai.tool.call.arguments, gen_ai.tool.call.result + if "tool_arguments" in attrs and "gen_ai.tool.call.arguments" not in attrs: + attrs["gen_ai.tool.call.arguments"] = attrs["tool_arguments"] + modified = True + + if "tool_response" in attrs and "gen_ai.tool.call.result" not in attrs: + attrs["gen_ai.tool.call.result"] = attrs["tool_response"] + modified = True + + if modified: + return _SpanWithModifiedAttributes(span, attrs) + + return span + + def _normalize_messages(self, messages: Any) -> Any: + """ + Normalize messages from Pydantic AI format to OpenAI chat format. + + Pydantic AI: {"parts": [{"content": "...", "type": "text"}], "role": "user"} + OpenAI: {"content": "...", "role": "user"} + """ + import json + + # Parse if string + if isinstance(messages, str): + try: + messages = json.loads(messages) + except json.JSONDecodeError: + return messages + + if not isinstance(messages, list): + return messages + + normalized: List[Dict[str, Any]] = [] + + for msg in messages: + if not isinstance(msg, dict): + normalized.append(msg) + continue + + # Check if this is Pydantic AI format with "parts" + if "parts" in msg and isinstance(msg["parts"], list): + normalized_msg = self._normalize_pydantic_message(msg) + normalized.append(normalized_msg) + else: + # Already in standard format + normalized.append(msg) + + return normalized + + def _normalize_pydantic_message(self, msg: Dict[str, Any]) -> Dict[str, Any]: + """Convert a single Pydantic AI message to OpenAI format.""" + role = msg.get("role", "unknown") + parts = msg.get("parts", []) + + text_parts: List[str] = [] + tool_calls: List[Dict[str, Any]] = [] + + for part in parts: + if not isinstance(part, dict): + continue + + part_type = part.get("type", "text") + + if part_type == "text" and "content" in part: + text_parts.append(str(part["content"])) + elif part_type == "tool_call": + tool_calls.append({ + "id": part.get("id", ""), + "type": "function", + "function": { + "name": part.get("name", ""), + "arguments": part.get("arguments", "{}"), + } + }) + + # Build normalized message + normalized: Dict[str, Any] = {"role": role} + + if text_parts: + normalized["content"] = "\n".join(text_parts) if len(text_parts) > 1 else text_parts[0] + elif not tool_calls: + normalized["content"] = "" + + if tool_calls: + normalized["tool_calls"] = tool_calls + + # Preserve finish_reason if present (for output/assistant messages) + if "finish_reason" in msg: + normalized["finish_reason"] = msg["finish_reason"] + + return normalized + + +class _SpanWithModifiedAttributes: + """ + Wrapper that presents a span with modified attributes. + + This allows us to transform attributes without mutating the original span. + """ + + def __init__(self, original_span: ReadableSpan, modified_attrs: Dict[str, Any]): + self._original = original_span + self._modified_attrs = modified_attrs + + @property + def attributes(self) -> Dict[str, Any]: + return self._modified_attrs + + @property + def name(self) -> str: + return self._original.name + + @property + def context(self): + return self._original.context + + @property + def parent(self): + return self._original.parent + + @property + def start_time(self): + return self._original.start_time + + @property + def end_time(self): + return self._original.end_time + + @property + def status(self): + return self._original.status + + @property + def events(self): + return self._original.events + + @property + def links(self): + return self._original.links + + @property + def resource(self): + return self._original.resource + + @property + def instrumentation_scope(self): + return self._original.instrumentation_scope diff --git a/posthog/ai/pydantic_ai/instrument.py b/posthog/ai/pydantic_ai/instrument.py new file mode 100644 index 00000000..7aa254cf --- /dev/null +++ b/posthog/ai/pydantic_ai/instrument.py @@ -0,0 +1,92 @@ +""" +Pydantic AI instrumentation for PostHog. + +Provides a simple one-liner to instrument all Pydantic AI agents with PostHog tracing. +""" + +from typing import Any, Dict, Optional + +from posthog.client import Client as PostHogClient + + +def instrument_pydantic_ai( + client: PostHogClient, + distinct_id: Optional[str] = None, + privacy_mode: bool = False, + properties: Optional[Dict[str, Any]] = None, + groups: Optional[Dict[str, Any]] = None, + debug: bool = False, +) -> None: + """ + Instrument all Pydantic AI agents with PostHog tracing. + + This function sets up OpenTelemetry tracing for Pydantic AI and routes + all spans to PostHog as AI events ($ai_generation, $ai_trace, $ai_span). + + Usage: + from posthog import Posthog + from posthog.ai.pydantic_ai import instrument_pydantic_ai + from pydantic_ai import Agent + + posthog = Posthog(api_key="...", host="...") + instrument_pydantic_ai(posthog, distinct_id="user_123") + + # Now use Pydantic AI normally - all traces go to PostHog + agent = Agent('openai:gpt-4') + result = await agent.run('Hello!') + + Args: + client: PostHog client instance for sending events + distinct_id: Default distinct ID for all events. If not provided, + events will use the trace ID as distinct_id. + privacy_mode: If True, message content will be redacted from events. + Useful for sensitive data. + properties: Additional properties to include in all events + groups: PostHog groups to associate with all events + debug: Enable debug logging for troubleshooting + + Raises: + ImportError: If pydantic-ai or opentelemetry-sdk is not installed + """ + try: + from pydantic_ai import Agent + from pydantic_ai.models.instrumented import InstrumentationSettings + except ImportError as e: + raise ImportError( + "pydantic-ai is required for Pydantic AI instrumentation. " + "Install it with: pip install pydantic-ai" + ) from e + + try: + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import BatchSpanProcessor + except ImportError as e: + raise ImportError( + "opentelemetry-sdk is required for Pydantic AI instrumentation. " + "Install it with: pip install opentelemetry-sdk" + ) from e + + from posthog.ai.pydantic_ai.exporter import PydanticAISpanExporter + + # Create the Pydantic AI-specific exporter (handles message format normalization) + exporter = PydanticAISpanExporter( + client=client, + distinct_id=distinct_id, + privacy_mode=privacy_mode, + properties=properties, + groups=groups, + debug=debug, + ) + + # Create a TracerProvider with our exporter + provider = TracerProvider() + provider.add_span_processor(BatchSpanProcessor(exporter)) + + # Configure Pydantic AI instrumentation settings + settings = InstrumentationSettings( + tracer_provider=provider, + include_content=not privacy_mode, + ) + + # Apply instrumentation globally to all agents + Agent.instrument_all(settings) diff --git a/posthog/test/ai/otel/__init__.py b/posthog/test/ai/otel/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/posthog/test/ai/otel/test_exporter.py b/posthog/test/ai/otel/test_exporter.py new file mode 100644 index 00000000..d6f9fd77 --- /dev/null +++ b/posthog/test/ai/otel/test_exporter.py @@ -0,0 +1,440 @@ +""" +Tests for PostHogSpanExporter - the generic OpenTelemetry span exporter for PostHog. +""" + +import json +from unittest.mock import MagicMock, patch + +import pytest + +try: + from opentelemetry.sdk.trace import ReadableSpan + from opentelemetry.sdk.trace.export import SpanExportResult + from opentelemetry.trace import SpanContext, Status, StatusCode, TraceFlags + + from posthog.ai.otel import PostHogSpanExporter + + OTEL_AVAILABLE = True +except ImportError: + OTEL_AVAILABLE = False + +pytestmark = pytest.mark.skipif( + not OTEL_AVAILABLE, reason="OpenTelemetry SDK is not available" +) + + +@pytest.fixture +def mock_client(): + client = MagicMock() + client.privacy_mode = False + return client + + +def create_mock_span( + name: str, + attributes: dict = None, + trace_id: int = 0x418BB9C71D1C0591CD2AD7F97B58B9EB, + span_id: int = 0x1234567890ABCDEF, + parent_span_id: int = None, + start_time: int = 1000000000, + end_time: int = 2000000000, + status_code: StatusCode = StatusCode.OK, + status_description: str = None, +): + """Create a mock ReadableSpan for testing.""" + span = MagicMock(spec=ReadableSpan) + span.name = name + span.attributes = attributes or {} + + # Set up span context + span.context = MagicMock() + span.context.trace_id = trace_id + span.context.span_id = span_id + + # Set up parent context + if parent_span_id: + span.parent = MagicMock() + span.parent.span_id = parent_span_id + else: + span.parent = None + + span.start_time = start_time + span.end_time = end_time + + # Set up status + span.status = MagicMock() + span.status.status_code = status_code + span.status.description = status_description + + return span + + +class TestTraceIdFormatting: + """Tests for trace ID formatting to UUID format.""" + + def test_format_trace_id_as_uuid(self, mock_client): + exporter = PostHogSpanExporter(mock_client) + trace_id = 0x418BB9C71D1C0591CD2AD7F97B58B9EB + result = exporter._format_trace_id_as_uuid(trace_id) + assert result == "418bb9c7-1d1c-0591-cd2a-d7f97b58b9eb" + + def test_format_trace_id_preserves_leading_zeros(self, mock_client): + exporter = PostHogSpanExporter(mock_client) + trace_id = 0x00000000000000000000000000000001 + result = exporter._format_trace_id_as_uuid(trace_id) + assert result == "00000000-0000-0000-0000-000000000001" + + +class TestSpanClassification: + """Tests for span type classification logic.""" + + def test_is_generation_span_chat_prefix(self, mock_client): + exporter = PostHogSpanExporter(mock_client) + assert exporter._is_generation_span("chat openai", {}) is True + assert exporter._is_generation_span("chat gpt-4", {}) is True + + def test_is_generation_span_with_operation_name(self, mock_client): + exporter = PostHogSpanExporter(mock_client) + attrs = {"gen_ai.operation.name": "chat"} + assert exporter._is_generation_span("some_span", attrs) is True + + def test_is_generation_span_with_request_model(self, mock_client): + exporter = PostHogSpanExporter(mock_client) + attrs = {"gen_ai.request.model": "gpt-4"} + assert exporter._is_generation_span("some_span", attrs) is True + + def test_is_generation_span_negative(self, mock_client): + exporter = PostHogSpanExporter(mock_client) + assert exporter._is_generation_span("tool call", {}) is False + + def test_is_agent_span(self, mock_client): + exporter = PostHogSpanExporter(mock_client) + assert exporter._is_agent_span("agent run", {}) is True + assert exporter._is_agent_span("invoke_agent", {}) is True + assert exporter._is_agent_span("some_span", {"gen_ai.agent.name": "test"}) # truthy + assert not exporter._is_agent_span("some_span", {}) + + def test_is_tool_span(self, mock_client): + exporter = PostHogSpanExporter(mock_client) + assert exporter._is_tool_span("execute_tool get_weather", {}) is True + assert exporter._is_tool_span("running tools", {}) is True + assert exporter._is_tool_span("some_span", {"gen_ai.tool.name": "get_weather"}) is True + assert exporter._is_tool_span("model call", {}) is False + + +class TestGenerationEventCreation: + """Tests for $ai_generation event creation from model request spans.""" + + def test_basic_generation_event(self, mock_client): + exporter = PostHogSpanExporter(mock_client, distinct_id="user_123") + + span = create_mock_span( + name="chat openai", + attributes={ + "gen_ai.request.model": "gpt-4", + "gen_ai.system": "openai", + "gen_ai.usage.input_tokens": 100, + "gen_ai.usage.output_tokens": 50, + "gen_ai.input.messages": json.dumps([{"role": "user", "content": "Hello"}]), + "gen_ai.output.messages": json.dumps([{"role": "assistant", "content": "Hi!"}]), + }, + ) + + exporter.export([span]) + + mock_client.capture.assert_called_once() + call_kwargs = mock_client.capture.call_args[1] + + assert call_kwargs["event"] == "$ai_generation" + assert call_kwargs["distinct_id"] == "user_123" + + props = call_kwargs["properties"] + assert props["$ai_model"] == "gpt-4" + assert props["$ai_provider"] == "openai" + assert props["$ai_input_tokens"] == 100 + assert props["$ai_output_tokens"] == 50 + assert props["$ai_input"] == [{"role": "user", "content": "Hello"}] + assert props["$ai_output_choices"] == [{"role": "assistant", "content": "Hi!"}] + assert props["$ai_is_error"] is False + assert props["$ai_http_status"] == 200 + assert "$ai_trace_id" in props + assert "-" in props["$ai_trace_id"] # UUID format + + def test_generation_event_with_error(self, mock_client): + exporter = PostHogSpanExporter(mock_client, distinct_id="user_123") + + span = create_mock_span( + name="chat openai", + attributes={"gen_ai.request.model": "gpt-4"}, + status_code=StatusCode.ERROR, + status_description="Rate limit exceeded", + ) + + exporter.export([span]) + + props = mock_client.capture.call_args[1]["properties"] + assert props["$ai_is_error"] is True + assert props["$ai_http_status"] == 500 + assert props["$ai_error"] == "Rate limit exceeded" + + def test_generation_event_privacy_mode(self, mock_client): + exporter = PostHogSpanExporter(mock_client, distinct_id="user_123", privacy_mode=True) + + span = create_mock_span( + name="chat openai", + attributes={ + "gen_ai.request.model": "gpt-4", + "gen_ai.input.messages": json.dumps([{"role": "user", "content": "Secret data"}]), + "gen_ai.output.messages": json.dumps([{"role": "assistant", "content": "Response"}]), + }, + ) + + exporter.export([span]) + + props = mock_client.capture.call_args[1]["properties"] + assert "$ai_input" not in props + assert "$ai_output_choices" not in props + + def test_generation_event_with_model_parameters(self, mock_client): + exporter = PostHogSpanExporter(mock_client, distinct_id="user_123") + + span = create_mock_span( + name="chat openai", + attributes={ + "gen_ai.request.model": "gpt-4", + "gen_ai.request.temperature": 0.7, + "gen_ai.request.max_tokens": 1000, + "gen_ai.request.top_p": 0.9, + }, + ) + + exporter.export([span]) + + props = mock_client.capture.call_args[1]["properties"] + assert props["$ai_model_parameters"] == { + "temperature": 0.7, + "max_tokens": 1000, + "top_p": 0.9, + } + + +class TestAgentSpanHandling: + """Tests for agent span handling (should be skipped).""" + + def test_agent_span_is_skipped(self, mock_client): + exporter = PostHogSpanExporter(mock_client, distinct_id="user_123") + + span = create_mock_span(name="agent run", attributes={"gen_ai.agent.name": "TestAgent"}) + + exporter.export([span]) + + mock_client.capture.assert_not_called() + + def test_invoke_agent_span_is_skipped(self, mock_client): + exporter = PostHogSpanExporter(mock_client, distinct_id="user_123") + + span = create_mock_span(name="invoke_agent", attributes={}) + + exporter.export([span]) + + mock_client.capture.assert_not_called() + + +class TestToolSpanEventCreation: + """Tests for $ai_span event creation from tool execution spans.""" + + def test_basic_tool_span_event(self, mock_client): + exporter = PostHogSpanExporter(mock_client, distinct_id="user_123") + + span = create_mock_span( + name="execute_tool get_weather", + attributes={ + "gen_ai.tool.name": "get_weather", + "gen_ai.tool.call.arguments": json.dumps({"latitude": 37.7749, "longitude": -122.4194}), + "gen_ai.tool.call.result": "Sunny, 72°F", + }, + parent_span_id=0xABCDEF1234567890, + ) + + exporter.export([span]) + + call_kwargs = mock_client.capture.call_args[1] + assert call_kwargs["event"] == "$ai_span" + + props = call_kwargs["properties"] + assert props["$ai_span_name"] == "get_weather" + assert "$ai_trace_id" in props + assert "$ai_span_id" in props + assert "$ai_parent_id" in props + assert props["$ai_tool_arguments"] == {"latitude": 37.7749, "longitude": -122.4194} + assert props["$ai_tool_result"] == "Sunny, 72°F" + + def test_tool_span_privacy_mode(self, mock_client): + exporter = PostHogSpanExporter(mock_client, distinct_id="user_123", privacy_mode=True) + + span = create_mock_span( + name="execute_tool get_weather", + attributes={ + "gen_ai.tool.name": "get_weather", + "gen_ai.tool.call.arguments": json.dumps({"secret": "value"}), + "gen_ai.tool.call.result": "Secret result", + }, + ) + + exporter.export([span]) + + props = mock_client.capture.call_args[1]["properties"] + assert "$ai_tool_arguments" not in props + assert "$ai_tool_result" not in props + + +class TestDistinctIdHandling: + """Tests for distinct_id resolution.""" + + def test_distinct_id_from_constructor(self, mock_client): + exporter = PostHogSpanExporter(mock_client, distinct_id="configured_user") + + span = create_mock_span(name="chat openai", attributes={"gen_ai.request.model": "gpt-4"}) + + exporter.export([span]) + + assert mock_client.capture.call_args[1]["distinct_id"] == "configured_user" + + def test_distinct_id_from_span_attribute(self, mock_client): + exporter = PostHogSpanExporter(mock_client, distinct_id="default_user") + + span = create_mock_span( + name="chat openai", + attributes={ + "gen_ai.request.model": "gpt-4", + "posthog.distinct_id": "span_user", + }, + ) + + exporter.export([span]) + + assert mock_client.capture.call_args[1]["distinct_id"] == "span_user" + + def test_distinct_id_fallback_to_trace_id(self, mock_client): + exporter = PostHogSpanExporter(mock_client) + + span = create_mock_span( + name="chat openai", + attributes={"gen_ai.request.model": "gpt-4"}, + trace_id=0xABCDEF1234567890ABCDEF1234567890, + ) + + exporter.export([span]) + + assert mock_client.capture.call_args[1]["distinct_id"] == "abcdef1234567890abcdef1234567890" + + def test_process_person_profile_false_when_no_distinct_id(self, mock_client): + exporter = PostHogSpanExporter(mock_client) + + span = create_mock_span(name="chat openai", attributes={"gen_ai.request.model": "gpt-4"}) + + exporter.export([span]) + + props = mock_client.capture.call_args[1]["properties"] + assert props["$process_person_profile"] is False + + +class TestAdditionalProperties: + """Tests for additional properties and groups.""" + + def test_additional_properties_included(self, mock_client): + exporter = PostHogSpanExporter( + mock_client, + distinct_id="user_123", + properties={"$ai_session_id": "session_abc", "custom_prop": "value"}, + ) + + span = create_mock_span(name="chat openai", attributes={"gen_ai.request.model": "gpt-4"}) + + exporter.export([span]) + + props = mock_client.capture.call_args[1]["properties"] + assert props["$ai_session_id"] == "session_abc" + assert props["custom_prop"] == "value" + + def test_groups_included(self, mock_client): + exporter = PostHogSpanExporter( + mock_client, + distinct_id="user_123", + groups={"company": "posthog", "team": "product"}, + ) + + span = create_mock_span(name="chat openai", attributes={"gen_ai.request.model": "gpt-4"}) + + exporter.export([span]) + + assert mock_client.capture.call_args[1]["groups"] == { + "company": "posthog", + "team": "product", + } + + +class TestExportResult: + """Tests for export method return values.""" + + def test_export_returns_success(self, mock_client): + exporter = PostHogSpanExporter(mock_client, distinct_id="user_123") + + span = create_mock_span(name="chat openai", attributes={"gen_ai.request.model": "gpt-4"}) + + result = exporter.export([span]) + + assert result == SpanExportResult.SUCCESS + + def test_export_handles_exceptions_gracefully(self, mock_client): + mock_client.capture.side_effect = Exception("Network error") + exporter = PostHogSpanExporter(mock_client, distinct_id="user_123") + + span = create_mock_span(name="chat openai", attributes={"gen_ai.request.model": "gpt-4"}) + + result = exporter.export([span]) + + assert result == SpanExportResult.SUCCESS + + +class TestLatencyCalculation: + """Tests for latency calculation from span times.""" + + def test_latency_calculated_correctly(self, mock_client): + exporter = PostHogSpanExporter(mock_client, distinct_id="user_123") + + span = create_mock_span( + name="chat openai", + attributes={"gen_ai.request.model": "gpt-4"}, + start_time=1_000_000_000, # 1 second in nanoseconds + end_time=2_500_000_000, # 2.5 seconds in nanoseconds + ) + + exporter.export([span]) + + props = mock_client.capture.call_args[1]["properties"] + assert props["$ai_latency"] == 1.5 + + +class TestJsonParsing: + """Tests for JSON attribute parsing.""" + + def test_parse_json_string(self, mock_client): + exporter = PostHogSpanExporter(mock_client) + result = exporter._parse_json_attr('{"key": "value"}') + assert result == {"key": "value"} + + def test_parse_json_already_parsed(self, mock_client): + exporter = PostHogSpanExporter(mock_client) + result = exporter._parse_json_attr({"key": "value"}) + assert result == {"key": "value"} + + def test_parse_json_invalid_returns_original(self, mock_client): + exporter = PostHogSpanExporter(mock_client) + result = exporter._parse_json_attr("not json") + assert result == "not json" + + def test_parse_json_none(self, mock_client): + exporter = PostHogSpanExporter(mock_client) + result = exporter._parse_json_attr(None) + assert result is None diff --git a/posthog/test/ai/pydantic_ai/__init__.py b/posthog/test/ai/pydantic_ai/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/posthog/test/ai/pydantic_ai/test_exporter.py b/posthog/test/ai/pydantic_ai/test_exporter.py new file mode 100644 index 00000000..178da5ec --- /dev/null +++ b/posthog/test/ai/pydantic_ai/test_exporter.py @@ -0,0 +1,537 @@ +""" +Tests for PydanticAISpanExporter - handles Pydantic AI message format normalization. +""" + +import json +from unittest.mock import MagicMock, patch + +import pytest + +try: + from opentelemetry.sdk.trace import ReadableSpan + from opentelemetry.sdk.trace.export import SpanExportResult + from opentelemetry.trace import SpanContext, Status, StatusCode, TraceFlags + + from posthog.ai.pydantic_ai.exporter import PydanticAISpanExporter + + OTEL_AVAILABLE = True +except ImportError: + OTEL_AVAILABLE = False + +pytestmark = pytest.mark.skipif( + not OTEL_AVAILABLE, reason="OpenTelemetry SDK is not available" +) + + +@pytest.fixture +def mock_client(): + client = MagicMock() + client.privacy_mode = False + return client + + +def create_mock_span( + name: str, + attributes: dict = None, + trace_id: int = 0x418BB9C71D1C0591CD2AD7F97B58B9EB, + span_id: int = 0x1234567890ABCDEF, + parent_span_id: int = None, + start_time: int = 1000000000, + end_time: int = 2000000000, + status_code: StatusCode = StatusCode.OK, +): + """Create a mock ReadableSpan for testing.""" + span = MagicMock(spec=ReadableSpan) + span.name = name + span.attributes = attributes or {} + span.context = MagicMock() + span.context.trace_id = trace_id + span.context.span_id = span_id + if parent_span_id: + span.parent = MagicMock() + span.parent.span_id = parent_span_id + else: + span.parent = None + span.start_time = start_time + span.end_time = end_time + span.status = MagicMock() + span.status.status_code = status_code + span.status.description = None + return span + + +class TestMessageNormalization: + """Tests for normalizing Pydantic AI message format to OpenAI format.""" + + def test_normalize_simple_text_message(self, mock_client): + exporter = PydanticAISpanExporter(mock_client, distinct_id="user_123") + + pydantic_format = [ + {"parts": [{"content": "Hello, how are you?", "type": "text"}], "role": "user"} + ] + result = exporter._normalize_messages(pydantic_format) + + assert result == [{"content": "Hello, how are you?", "role": "user"}] + + def test_normalize_multiple_text_parts(self, mock_client): + exporter = PydanticAISpanExporter(mock_client, distinct_id="user_123") + + pydantic_format = [ + { + "parts": [ + {"content": "First part.", "type": "text"}, + {"content": "Second part.", "type": "text"}, + ], + "role": "user", + } + ] + result = exporter._normalize_messages(pydantic_format) + + assert result == [{"content": "First part.\nSecond part.", "role": "user"}] + + def test_normalize_message_with_tool_call(self, mock_client): + exporter = PydanticAISpanExporter(mock_client, distinct_id="user_123") + + pydantic_format = [ + { + "parts": [ + { + "type": "tool_call", + "id": "call_123", + "name": "get_weather", + "arguments": '{"latitude": 37.7749}', + } + ], + "role": "assistant", + } + ] + result = exporter._normalize_messages(pydantic_format) + + assert result == [ + { + "role": "assistant", + "tool_calls": [ + { + "id": "call_123", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"latitude": 37.7749}', + }, + } + ], + } + ] + + def test_normalize_message_with_text_and_tool_call(self, mock_client): + exporter = PydanticAISpanExporter(mock_client, distinct_id="user_123") + + pydantic_format = [ + { + "parts": [ + {"content": "Let me get the weather.", "type": "text"}, + { + "type": "tool_call", + "id": "call_123", + "name": "get_weather", + "arguments": "{}", + }, + ], + "role": "assistant", + } + ] + result = exporter._normalize_messages(pydantic_format) + + assert result == [ + { + "role": "assistant", + "content": "Let me get the weather.", + "tool_calls": [ + { + "id": "call_123", + "type": "function", + "function": {"name": "get_weather", "arguments": "{}"}, + } + ], + } + ] + + def test_normalize_preserves_finish_reason(self, mock_client): + exporter = PydanticAISpanExporter(mock_client, distinct_id="user_123") + + pydantic_format = [ + { + "parts": [{"content": "Done!", "type": "text"}], + "role": "assistant", + "finish_reason": "stop", + } + ] + result = exporter._normalize_messages(pydantic_format) + + assert result == [{"content": "Done!", "role": "assistant", "finish_reason": "stop"}] + + def test_normalize_already_openai_format(self, mock_client): + exporter = PydanticAISpanExporter(mock_client, distinct_id="user_123") + + openai_format = [ + {"content": "Hello!", "role": "user"}, + {"content": "Hi there!", "role": "assistant"}, + ] + result = exporter._normalize_messages(openai_format) + + assert result == openai_format + + def test_normalize_json_string_input(self, mock_client): + exporter = PydanticAISpanExporter(mock_client, distinct_id="user_123") + + json_string = json.dumps( + [{"parts": [{"content": "Hello", "type": "text"}], "role": "user"}] + ) + result = exporter._normalize_messages(json_string) + + assert result == [{"content": "Hello", "role": "user"}] + + def test_normalize_invalid_json_returns_original(self, mock_client): + exporter = PydanticAISpanExporter(mock_client, distinct_id="user_123") + + invalid_json = "not valid json" + result = exporter._normalize_messages(invalid_json) + + assert result == "not valid json" + + def test_normalize_empty_parts(self, mock_client): + exporter = PydanticAISpanExporter(mock_client, distinct_id="user_123") + + pydantic_format = [{"parts": [], "role": "user"}] + result = exporter._normalize_messages(pydantic_format) + + assert result == [{"role": "user", "content": ""}] + + +class TestSpanTransformation: + """Tests for span attribute transformation.""" + + def test_transform_span_normalizes_input_messages(self, mock_client): + exporter = PydanticAISpanExporter(mock_client, distinct_id="user_123") + + pydantic_messages = json.dumps( + [{"parts": [{"content": "Hello", "type": "text"}], "role": "user"}] + ) + + span = create_mock_span( + name="chat openai", + attributes={ + "gen_ai.request.model": "gpt-4", + "gen_ai.input.messages": pydantic_messages, + }, + ) + + transformed = exporter._transform_span(span) + + # Compare as parsed JSON to avoid key ordering issues + result = json.loads(transformed.attributes["gen_ai.input.messages"]) + assert result == [{"content": "Hello", "role": "user"}] + + def test_transform_span_normalizes_output_messages(self, mock_client): + exporter = PydanticAISpanExporter(mock_client, distinct_id="user_123") + + pydantic_messages = json.dumps( + [{"parts": [{"content": "Hi!", "type": "text"}], "role": "assistant"}] + ) + + span = create_mock_span( + name="chat openai", + attributes={ + "gen_ai.request.model": "gpt-4", + "gen_ai.output.messages": pydantic_messages, + }, + ) + + transformed = exporter._transform_span(span) + + # Compare as parsed JSON to avoid key ordering issues + result = json.loads(transformed.attributes["gen_ai.output.messages"]) + assert result == [{"content": "Hi!", "role": "assistant"}] + + def test_transform_span_preserves_other_attributes(self, mock_client): + exporter = PydanticAISpanExporter(mock_client, distinct_id="user_123") + + span = create_mock_span( + name="chat openai", + attributes={ + "gen_ai.request.model": "gpt-4", + "gen_ai.usage.input_tokens": 100, + "gen_ai.usage.output_tokens": 50, + "gen_ai.input.messages": json.dumps( + [{"parts": [{"content": "Hi", "type": "text"}], "role": "user"}] + ), + }, + ) + + transformed = exporter._transform_span(span) + + assert transformed.attributes["gen_ai.request.model"] == "gpt-4" + assert transformed.attributes["gen_ai.usage.input_tokens"] == 100 + assert transformed.attributes["gen_ai.usage.output_tokens"] == 50 + + def test_transform_span_no_modification_preserves_content(self, mock_client): + exporter = PydanticAISpanExporter(mock_client, distinct_id="user_123") + + openai_messages = json.dumps([{"content": "Hi", "role": "user"}]) + + span = create_mock_span( + name="chat openai", + attributes={ + "gen_ai.request.model": "gpt-4", + "gen_ai.input.messages": openai_messages, + }, + ) + + transformed = exporter._transform_span(span) + + # Content should be equivalent even if span wrapper is different + result = json.loads(transformed.attributes["gen_ai.input.messages"]) + assert result == [{"content": "Hi", "role": "user"}] + + +class TestEndToEndExport: + """Tests for full export flow with message normalization.""" + + def test_export_normalizes_and_captures(self, mock_client): + exporter = PydanticAISpanExporter(mock_client, distinct_id="user_123") + + pydantic_input = json.dumps( + [{"parts": [{"content": "What's the weather?", "type": "text"}], "role": "user"}] + ) + pydantic_output = json.dumps( + [ + { + "parts": [{"content": "It's sunny!", "type": "text"}], + "role": "assistant", + "finish_reason": "stop", + } + ] + ) + + span = create_mock_span( + name="chat openai", + attributes={ + "gen_ai.request.model": "gpt-4", + "gen_ai.system": "openai", + "gen_ai.input.messages": pydantic_input, + "gen_ai.output.messages": pydantic_output, + }, + ) + + result = exporter.export([span]) + + assert result == SpanExportResult.SUCCESS + mock_client.capture.assert_called_once() + + props = mock_client.capture.call_args[1]["properties"] + assert props["$ai_input"] == [{"content": "What's the weather?", "role": "user"}] + assert props["$ai_output_choices"] == [ + {"content": "It's sunny!", "role": "assistant", "finish_reason": "stop"} + ] + + def test_export_with_tool_calls_normalized(self, mock_client): + exporter = PydanticAISpanExporter(mock_client, distinct_id="user_123") + + pydantic_output = json.dumps( + [ + { + "parts": [ + { + "type": "tool_call", + "id": "call_abc", + "name": "get_weather", + "arguments": '{"city": "SF"}', + } + ], + "role": "assistant", + } + ] + ) + + span = create_mock_span( + name="chat openai", + attributes={ + "gen_ai.request.model": "gpt-4", + "gen_ai.output.messages": pydantic_output, + }, + ) + + exporter.export([span]) + + props = mock_client.capture.call_args[1]["properties"] + expected_tool_call = { + "role": "assistant", + "tool_calls": [ + { + "id": "call_abc", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "SF"}'}, + } + ], + } + assert props["$ai_output_choices"] == [expected_tool_call] + + +class TestToolAttributeMapping: + """Tests for mapping Pydantic AI tool attributes to GenAI standard names.""" + + def test_maps_tool_arguments(self, mock_client): + exporter = PydanticAISpanExporter(mock_client, distinct_id="user_123") + + span = create_mock_span( + name="running tool get_weather", + attributes={ + "tool_arguments": '{"city": "SF"}', + }, + ) + + transformed = exporter._transform_span(span) + + assert transformed.attributes["gen_ai.tool.call.arguments"] == '{"city": "SF"}' + assert transformed.attributes["tool_arguments"] == '{"city": "SF"}' + + def test_maps_tool_response(self, mock_client): + exporter = PydanticAISpanExporter(mock_client, distinct_id="user_123") + + span = create_mock_span( + name="running tool get_weather", + attributes={ + "tool_response": "Sunny, 72°F", + }, + ) + + transformed = exporter._transform_span(span) + + assert transformed.attributes["gen_ai.tool.call.result"] == "Sunny, 72°F" + assert transformed.attributes["tool_response"] == "Sunny, 72°F" + + def test_maps_both_tool_attributes(self, mock_client): + exporter = PydanticAISpanExporter(mock_client, distinct_id="user_123") + + span = create_mock_span( + name="running tool get_weather", + attributes={ + "tool_arguments": '{"city": "SF"}', + "tool_response": "Sunny, 72°F", + }, + ) + + transformed = exporter._transform_span(span) + + assert transformed.attributes["gen_ai.tool.call.arguments"] == '{"city": "SF"}' + assert transformed.attributes["gen_ai.tool.call.result"] == "Sunny, 72°F" + + def test_does_not_overwrite_existing_genai_attributes(self, mock_client): + exporter = PydanticAISpanExporter(mock_client, distinct_id="user_123") + + span = create_mock_span( + name="running tool get_weather", + attributes={ + "tool_arguments": '{"city": "SF"}', + "gen_ai.tool.call.arguments": '{"existing": "value"}', + }, + ) + + transformed = exporter._transform_span(span) + + # Should preserve existing GenAI attribute, not overwrite + assert transformed.attributes["gen_ai.tool.call.arguments"] == '{"existing": "value"}' + + def test_tool_span_export_with_mapped_attributes(self, mock_client): + exporter = PydanticAISpanExporter(mock_client, distinct_id="user_123") + + span = create_mock_span( + name="running tool get_weather", + attributes={ + "gen_ai.tool.name": "get_weather", + "tool_arguments": '{"city": "SF"}', + "tool_response": "Sunny, 72°F", + }, + parent_span_id=0xABCDEF1234567890, + ) + + exporter.export([span]) + + props = mock_client.capture.call_args[1]["properties"] + assert props["$ai_tool_arguments"] == {"city": "SF"} + assert props["$ai_tool_result"] == "Sunny, 72°F" + + +class TestSpanWrapperProperties: + """Tests for the _SpanWithModifiedAttributes wrapper.""" + + def test_wrapper_preserves_span_name(self, mock_client): + exporter = PydanticAISpanExporter(mock_client, distinct_id="user_123") + + span = create_mock_span( + name="chat openai", + attributes={ + "gen_ai.input.messages": json.dumps( + [{"parts": [{"content": "Hi", "type": "text"}], "role": "user"}] + ) + }, + ) + + transformed = exporter._transform_span(span) + + assert transformed.name == "chat openai" + + def test_wrapper_preserves_context(self, mock_client): + exporter = PydanticAISpanExporter(mock_client, distinct_id="user_123") + + span = create_mock_span( + name="chat openai", + trace_id=0x12345, + span_id=0xABCDE, + attributes={ + "gen_ai.input.messages": json.dumps( + [{"parts": [{"content": "Hi", "type": "text"}], "role": "user"}] + ) + }, + ) + + transformed = exporter._transform_span(span) + + assert transformed.context.trace_id == 0x12345 + assert transformed.context.span_id == 0xABCDE + + def test_wrapper_preserves_timing(self, mock_client): + exporter = PydanticAISpanExporter(mock_client, distinct_id="user_123") + + span = create_mock_span( + name="chat openai", + start_time=1000, + end_time=2000, + attributes={ + "gen_ai.input.messages": json.dumps( + [{"parts": [{"content": "Hi", "type": "text"}], "role": "user"}] + ) + }, + ) + + transformed = exporter._transform_span(span) + + assert transformed.start_time == 1000 + assert transformed.end_time == 2000 + + def test_wrapper_preserves_status(self, mock_client): + exporter = PydanticAISpanExporter(mock_client, distinct_id="user_123") + + span = create_mock_span( + name="chat openai", + status_code=StatusCode.ERROR, + attributes={ + "gen_ai.input.messages": json.dumps( + [{"parts": [{"content": "Hi", "type": "text"}], "role": "user"}] + ) + }, + ) + + transformed = exporter._transform_span(span) + + assert transformed.status.status_code == StatusCode.ERROR diff --git a/posthog/test/ai/pydantic_ai/test_instrument.py b/posthog/test/ai/pydantic_ai/test_instrument.py new file mode 100644 index 00000000..8b164187 --- /dev/null +++ b/posthog/test/ai/pydantic_ai/test_instrument.py @@ -0,0 +1,97 @@ +""" +Tests for instrument_pydantic_ai function. +""" + +from unittest.mock import MagicMock, patch + +import pytest + +try: + from pydantic_ai import Agent + from pydantic_ai.models.instrumented import InstrumentationSettings + + DEPS_AVAILABLE = True +except ImportError: + DEPS_AVAILABLE = False + Agent = None + InstrumentationSettings = None + +pytestmark = pytest.mark.skipif( + not DEPS_AVAILABLE, reason="pydantic-ai and opentelemetry-sdk are required" +) + + +@pytest.fixture +def mock_client(): + client = MagicMock() + client.privacy_mode = False + return client + + +class TestInstrumentPydanticAI: + """Tests for the instrument_pydantic_ai function.""" + + def test_basic_instrumentation(self, mock_client): + from posthog.ai.pydantic_ai import instrument_pydantic_ai + + with patch.object(Agent, "instrument_all") as mock_instrument_all: + instrument_pydantic_ai(mock_client, distinct_id="user_123") + + mock_instrument_all.assert_called_once() + settings = mock_instrument_all.call_args[0][0] + assert isinstance(settings, InstrumentationSettings) + + def test_privacy_mode_disables_content(self, mock_client): + from posthog.ai.pydantic_ai import instrument_pydantic_ai + + with patch.object(Agent, "instrument_all") as mock_instrument_all: + instrument_pydantic_ai(mock_client, privacy_mode=True) + + settings = mock_instrument_all.call_args[0][0] + assert settings.include_content is False + + def test_privacy_mode_false_includes_content(self, mock_client): + from posthog.ai.pydantic_ai import instrument_pydantic_ai + + with patch.object(Agent, "instrument_all") as mock_instrument_all: + instrument_pydantic_ai(mock_client, privacy_mode=False) + + settings = mock_instrument_all.call_args[0][0] + assert settings.include_content is True + + def test_tracer_configured_via_settings(self, mock_client): + from posthog.ai.pydantic_ai import instrument_pydantic_ai + + with patch.object(Agent, "instrument_all") as mock_instrument_all: + instrument_pydantic_ai(mock_client, distinct_id="user_123") + + settings = mock_instrument_all.call_args[0][0] + # InstrumentationSettings creates a tracer internally from the provider + # We verify it's properly configured by checking it has a tracer attribute + assert hasattr(settings, "tracer") + + def test_accepts_properties(self, mock_client): + from posthog.ai.pydantic_ai import instrument_pydantic_ai + + with patch.object(Agent, "instrument_all") as mock_instrument_all: + properties = {"$ai_session_id": "session_123", "custom": "value"} + instrument_pydantic_ai(mock_client, properties=properties) + + mock_instrument_all.assert_called_once() + + def test_accepts_groups(self, mock_client): + from posthog.ai.pydantic_ai import instrument_pydantic_ai + + with patch.object(Agent, "instrument_all") as mock_instrument_all: + groups = {"company": "posthog", "team": "product"} + instrument_pydantic_ai(mock_client, groups=groups) + + mock_instrument_all.assert_called_once() + + def test_debug_mode(self, mock_client): + from posthog.ai.pydantic_ai import instrument_pydantic_ai + + with patch.object(Agent, "instrument_all") as mock_instrument_all: + instrument_pydantic_ai(mock_client, debug=True) + + mock_instrument_all.assert_called_once() diff --git a/pyproject.toml b/pyproject.toml index b0e69264..4cfaaf26 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,8 @@ Repository = "https://github.com/posthog/posthog-python" [project.optional-dependencies] langchain = ["langchain>=0.2.0"] +pydantic-ai = ["pydantic-ai>=0.2.0", "opentelemetry-sdk>=1.20.0"] +otel = ["opentelemetry-sdk>=1.20.0"] dev = [ "django-stubs", "lxml", @@ -75,6 +77,8 @@ test = [ "langchain-anthropic>=1.0", "google-genai", "pydantic", + "pydantic-ai>=0.2.0", + "opentelemetry-sdk>=1.20.0", "parameterized>=0.8.1", ] @@ -86,6 +90,8 @@ packages = [ "posthog.ai.openai", "posthog.ai.anthropic", "posthog.ai.gemini", + "posthog.ai.otel", + "posthog.ai.pydantic_ai", "posthog.test", "posthog.integrations", ] From 94922d5817e9c15b22bc2871b0d7f6ad23a02c0a Mon Sep 17 00:00:00 2001 From: Andrew Maguire Date: Tue, 2 Dec 2025 16:29:38 +0000 Subject: [PATCH 02/14] Update posthog/ai/otel/exporter.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- posthog/ai/otel/exporter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/posthog/ai/otel/exporter.py b/posthog/ai/otel/exporter.py index de12666a..144f6628 100644 --- a/posthog/ai/otel/exporter.py +++ b/posthog/ai/otel/exporter.py @@ -365,7 +365,7 @@ def _create_generation_event( "$ai_latency": latency, "$ai_http_status": 500 if is_error else 200, "$ai_is_error": is_error, - "$ai_framework": "pydantic-ai", + "$ai_framework": "opentelemetry", **self._properties, } From 807a9658643b5251c57512455865ab1b8c5a0719 Mon Sep 17 00:00:00 2001 From: Andrew Maguire Date: Tue, 2 Dec 2025 16:29:52 +0000 Subject: [PATCH 03/14] Update posthog/ai/otel/exporter.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- posthog/ai/otel/exporter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/posthog/ai/otel/exporter.py b/posthog/ai/otel/exporter.py index 144f6628..5bdd9901 100644 --- a/posthog/ai/otel/exporter.py +++ b/posthog/ai/otel/exporter.py @@ -449,7 +449,7 @@ def _create_tool_span_event( "$ai_span_name": tool_name, "$ai_latency": latency, "$ai_is_error": is_error, - "$ai_framework": "pydantic-ai", + "$ai_framework": "opentelemetry", **self._properties, } From 7f6821464cc2c1b4eb3d11d02339f99b0db82b66 Mon Sep 17 00:00:00 2001 From: Andrew Maguire Date: Tue, 2 Dec 2025 16:30:02 +0000 Subject: [PATCH 04/14] Update posthog/ai/otel/exporter.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- posthog/ai/otel/exporter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/posthog/ai/otel/exporter.py b/posthog/ai/otel/exporter.py index 5bdd9901..27ce24c8 100644 --- a/posthog/ai/otel/exporter.py +++ b/posthog/ai/otel/exporter.py @@ -492,7 +492,7 @@ def _create_span_event( "$ai_span_name": span.name, "$ai_latency": latency, "$ai_is_error": is_error, - "$ai_framework": "pydantic-ai", + "$ai_framework": "opentelemetry", **self._properties, } From 4da94db770d36819f41a0e7ecf9c7a67b5fd21a5 Mon Sep 17 00:00:00 2001 From: Andrew Maguire Date: Tue, 2 Dec 2025 16:30:57 +0000 Subject: [PATCH 05/14] style: Format files with ruff --- posthog/ai/otel/exporter.py | 35 ++++++--- posthog/ai/pydantic_ai/exporter.py | 28 +++++--- posthog/test/ai/otel/test_exporter.py | 75 +++++++++++++++----- posthog/test/ai/pydantic_ai/test_exporter.py | 25 +++++-- 4 files changed, 121 insertions(+), 42 deletions(-) diff --git a/posthog/ai/otel/exporter.py b/posthog/ai/otel/exporter.py index 27ce24c8..8bbdde7c 100644 --- a/posthog/ai/otel/exporter.py +++ b/posthog/ai/otel/exporter.py @@ -212,9 +212,7 @@ def _span_to_event(self, span: ReadableSpan) -> Optional[Dict[str, Any]]: trace_id = self._format_trace_id_as_uuid(span.context.trace_id) # Span IDs remain as hex (no dashes needed) span_id = format(span.context.span_id, "016x") - parent_span_id = ( - format(span.parent.span_id, "016x") if span.parent else None - ) + parent_span_id = format(span.parent.span_id, "016x") if span.parent else None # Check for error status is_error = span.status.status_code == StatusCode.ERROR if span.status else False @@ -223,7 +221,14 @@ def _span_to_event(self, span: ReadableSpan) -> Optional[Dict[str, Any]]: # Model request span → $ai_generation if self._is_generation_span(span_name, attrs): return self._create_generation_event( - span, attrs, trace_id, span_id, parent_span_id, latency, is_error, error_message + span, + attrs, + trace_id, + span_id, + parent_span_id, + latency, + is_error, + error_message, ) # Agent run span → skip (PostHog UI auto-creates trace wrapper from generation events) @@ -234,13 +239,27 @@ def _span_to_event(self, span: ReadableSpan) -> Optional[Dict[str, Any]]: # Tool execution span → $ai_span if self._is_tool_span(span_name, attrs): return self._create_tool_span_event( - span, attrs, trace_id, span_id, parent_span_id, latency, is_error, error_message + span, + attrs, + trace_id, + span_id, + parent_span_id, + latency, + is_error, + error_message, ) # Generic span that might be part of AI workflow if self._is_ai_related_span(span_name, attrs): return self._create_span_event( - span, attrs, trace_id, span_id, parent_span_id, latency, is_error, error_message + span, + attrs, + trace_id, + span_id, + parent_span_id, + latency, + is_error, + error_message, ) return None @@ -507,9 +526,7 @@ def _create_span_event( return {"name": "$ai_span", "properties": properties} - def _parse_json_attr( - self, value: Optional[Union[str, Any]] - ) -> Optional[Any]: + def _parse_json_attr(self, value: Optional[Union[str, Any]]) -> Optional[Any]: """Parse a JSON string attribute, returning the value as-is if already parsed.""" if value is None: return None diff --git a/posthog/ai/pydantic_ai/exporter.py b/posthog/ai/pydantic_ai/exporter.py index 85095c11..a5e56bdd 100644 --- a/posthog/ai/pydantic_ai/exporter.py +++ b/posthog/ai/pydantic_ai/exporter.py @@ -93,7 +93,9 @@ def _transform_span(self, span: ReadableSpan) -> ReadableSpan: normalized = self._normalize_messages(input_msgs) if normalized != input_msgs: attrs["gen_ai.input.messages"] = ( - json.dumps(normalized) if isinstance(normalized, list) else normalized + json.dumps(normalized) + if isinstance(normalized, list) + else normalized ) modified = True @@ -103,7 +105,9 @@ def _transform_span(self, span: ReadableSpan) -> ReadableSpan: normalized = self._normalize_messages(output_msgs) if normalized != output_msgs: attrs["gen_ai.output.messages"] = ( - json.dumps(normalized) if isinstance(normalized, list) else normalized + json.dumps(normalized) + if isinstance(normalized, list) + else normalized ) modified = True @@ -176,20 +180,24 @@ def _normalize_pydantic_message(self, msg: Dict[str, Any]) -> Dict[str, Any]: if part_type == "text" and "content" in part: text_parts.append(str(part["content"])) elif part_type == "tool_call": - tool_calls.append({ - "id": part.get("id", ""), - "type": "function", - "function": { - "name": part.get("name", ""), - "arguments": part.get("arguments", "{}"), + tool_calls.append( + { + "id": part.get("id", ""), + "type": "function", + "function": { + "name": part.get("name", ""), + "arguments": part.get("arguments", "{}"), + }, } - }) + ) # Build normalized message normalized: Dict[str, Any] = {"role": role} if text_parts: - normalized["content"] = "\n".join(text_parts) if len(text_parts) > 1 else text_parts[0] + normalized["content"] = ( + "\n".join(text_parts) if len(text_parts) > 1 else text_parts[0] + ) elif not tool_calls: normalized["content"] = "" diff --git a/posthog/test/ai/otel/test_exporter.py b/posthog/test/ai/otel/test_exporter.py index d6f9fd77..819dc291 100644 --- a/posthog/test/ai/otel/test_exporter.py +++ b/posthog/test/ai/otel/test_exporter.py @@ -111,14 +111,19 @@ def test_is_agent_span(self, mock_client): exporter = PostHogSpanExporter(mock_client) assert exporter._is_agent_span("agent run", {}) is True assert exporter._is_agent_span("invoke_agent", {}) is True - assert exporter._is_agent_span("some_span", {"gen_ai.agent.name": "test"}) # truthy + assert exporter._is_agent_span( + "some_span", {"gen_ai.agent.name": "test"} + ) # truthy assert not exporter._is_agent_span("some_span", {}) def test_is_tool_span(self, mock_client): exporter = PostHogSpanExporter(mock_client) assert exporter._is_tool_span("execute_tool get_weather", {}) is True assert exporter._is_tool_span("running tools", {}) is True - assert exporter._is_tool_span("some_span", {"gen_ai.tool.name": "get_weather"}) is True + assert ( + exporter._is_tool_span("some_span", {"gen_ai.tool.name": "get_weather"}) + is True + ) assert exporter._is_tool_span("model call", {}) is False @@ -135,8 +140,12 @@ def test_basic_generation_event(self, mock_client): "gen_ai.system": "openai", "gen_ai.usage.input_tokens": 100, "gen_ai.usage.output_tokens": 50, - "gen_ai.input.messages": json.dumps([{"role": "user", "content": "Hello"}]), - "gen_ai.output.messages": json.dumps([{"role": "assistant", "content": "Hi!"}]), + "gen_ai.input.messages": json.dumps( + [{"role": "user", "content": "Hello"}] + ), + "gen_ai.output.messages": json.dumps( + [{"role": "assistant", "content": "Hi!"}] + ), }, ) @@ -178,14 +187,20 @@ def test_generation_event_with_error(self, mock_client): assert props["$ai_error"] == "Rate limit exceeded" def test_generation_event_privacy_mode(self, mock_client): - exporter = PostHogSpanExporter(mock_client, distinct_id="user_123", privacy_mode=True) + exporter = PostHogSpanExporter( + mock_client, distinct_id="user_123", privacy_mode=True + ) span = create_mock_span( name="chat openai", attributes={ "gen_ai.request.model": "gpt-4", - "gen_ai.input.messages": json.dumps([{"role": "user", "content": "Secret data"}]), - "gen_ai.output.messages": json.dumps([{"role": "assistant", "content": "Response"}]), + "gen_ai.input.messages": json.dumps( + [{"role": "user", "content": "Secret data"}] + ), + "gen_ai.output.messages": json.dumps( + [{"role": "assistant", "content": "Response"}] + ), }, ) @@ -224,7 +239,9 @@ class TestAgentSpanHandling: def test_agent_span_is_skipped(self, mock_client): exporter = PostHogSpanExporter(mock_client, distinct_id="user_123") - span = create_mock_span(name="agent run", attributes={"gen_ai.agent.name": "TestAgent"}) + span = create_mock_span( + name="agent run", attributes={"gen_ai.agent.name": "TestAgent"} + ) exporter.export([span]) @@ -250,7 +267,9 @@ def test_basic_tool_span_event(self, mock_client): name="execute_tool get_weather", attributes={ "gen_ai.tool.name": "get_weather", - "gen_ai.tool.call.arguments": json.dumps({"latitude": 37.7749, "longitude": -122.4194}), + "gen_ai.tool.call.arguments": json.dumps( + {"latitude": 37.7749, "longitude": -122.4194} + ), "gen_ai.tool.call.result": "Sunny, 72°F", }, parent_span_id=0xABCDEF1234567890, @@ -266,11 +285,16 @@ def test_basic_tool_span_event(self, mock_client): assert "$ai_trace_id" in props assert "$ai_span_id" in props assert "$ai_parent_id" in props - assert props["$ai_tool_arguments"] == {"latitude": 37.7749, "longitude": -122.4194} + assert props["$ai_tool_arguments"] == { + "latitude": 37.7749, + "longitude": -122.4194, + } assert props["$ai_tool_result"] == "Sunny, 72°F" def test_tool_span_privacy_mode(self, mock_client): - exporter = PostHogSpanExporter(mock_client, distinct_id="user_123", privacy_mode=True) + exporter = PostHogSpanExporter( + mock_client, distinct_id="user_123", privacy_mode=True + ) span = create_mock_span( name="execute_tool get_weather", @@ -294,7 +318,9 @@ class TestDistinctIdHandling: def test_distinct_id_from_constructor(self, mock_client): exporter = PostHogSpanExporter(mock_client, distinct_id="configured_user") - span = create_mock_span(name="chat openai", attributes={"gen_ai.request.model": "gpt-4"}) + span = create_mock_span( + name="chat openai", attributes={"gen_ai.request.model": "gpt-4"} + ) exporter.export([span]) @@ -326,12 +352,17 @@ def test_distinct_id_fallback_to_trace_id(self, mock_client): exporter.export([span]) - assert mock_client.capture.call_args[1]["distinct_id"] == "abcdef1234567890abcdef1234567890" + assert ( + mock_client.capture.call_args[1]["distinct_id"] + == "abcdef1234567890abcdef1234567890" + ) def test_process_person_profile_false_when_no_distinct_id(self, mock_client): exporter = PostHogSpanExporter(mock_client) - span = create_mock_span(name="chat openai", attributes={"gen_ai.request.model": "gpt-4"}) + span = create_mock_span( + name="chat openai", attributes={"gen_ai.request.model": "gpt-4"} + ) exporter.export([span]) @@ -349,7 +380,9 @@ def test_additional_properties_included(self, mock_client): properties={"$ai_session_id": "session_abc", "custom_prop": "value"}, ) - span = create_mock_span(name="chat openai", attributes={"gen_ai.request.model": "gpt-4"}) + span = create_mock_span( + name="chat openai", attributes={"gen_ai.request.model": "gpt-4"} + ) exporter.export([span]) @@ -364,7 +397,9 @@ def test_groups_included(self, mock_client): groups={"company": "posthog", "team": "product"}, ) - span = create_mock_span(name="chat openai", attributes={"gen_ai.request.model": "gpt-4"}) + span = create_mock_span( + name="chat openai", attributes={"gen_ai.request.model": "gpt-4"} + ) exporter.export([span]) @@ -380,7 +415,9 @@ class TestExportResult: def test_export_returns_success(self, mock_client): exporter = PostHogSpanExporter(mock_client, distinct_id="user_123") - span = create_mock_span(name="chat openai", attributes={"gen_ai.request.model": "gpt-4"}) + span = create_mock_span( + name="chat openai", attributes={"gen_ai.request.model": "gpt-4"} + ) result = exporter.export([span]) @@ -390,7 +427,9 @@ def test_export_handles_exceptions_gracefully(self, mock_client): mock_client.capture.side_effect = Exception("Network error") exporter = PostHogSpanExporter(mock_client, distinct_id="user_123") - span = create_mock_span(name="chat openai", attributes={"gen_ai.request.model": "gpt-4"}) + span = create_mock_span( + name="chat openai", attributes={"gen_ai.request.model": "gpt-4"} + ) result = exporter.export([span]) diff --git a/posthog/test/ai/pydantic_ai/test_exporter.py b/posthog/test/ai/pydantic_ai/test_exporter.py index 178da5ec..6c7f5584 100644 --- a/posthog/test/ai/pydantic_ai/test_exporter.py +++ b/posthog/test/ai/pydantic_ai/test_exporter.py @@ -67,7 +67,10 @@ def test_normalize_simple_text_message(self, mock_client): exporter = PydanticAISpanExporter(mock_client, distinct_id="user_123") pydantic_format = [ - {"parts": [{"content": "Hello, how are you?", "type": "text"}], "role": "user"} + { + "parts": [{"content": "Hello, how are you?", "type": "text"}], + "role": "user", + } ] result = exporter._normalize_messages(pydantic_format) @@ -168,7 +171,9 @@ def test_normalize_preserves_finish_reason(self, mock_client): ] result = exporter._normalize_messages(pydantic_format) - assert result == [{"content": "Done!", "role": "assistant", "finish_reason": "stop"}] + assert result == [ + {"content": "Done!", "role": "assistant", "finish_reason": "stop"} + ] def test_normalize_already_openai_format(self, mock_client): exporter = PydanticAISpanExporter(mock_client, distinct_id="user_123") @@ -301,7 +306,12 @@ def test_export_normalizes_and_captures(self, mock_client): exporter = PydanticAISpanExporter(mock_client, distinct_id="user_123") pydantic_input = json.dumps( - [{"parts": [{"content": "What's the weather?", "type": "text"}], "role": "user"}] + [ + { + "parts": [{"content": "What's the weather?", "type": "text"}], + "role": "user", + } + ] ) pydantic_output = json.dumps( [ @@ -329,7 +339,9 @@ def test_export_normalizes_and_captures(self, mock_client): mock_client.capture.assert_called_once() props = mock_client.capture.call_args[1]["properties"] - assert props["$ai_input"] == [{"content": "What's the weather?", "role": "user"}] + assert props["$ai_input"] == [ + {"content": "What's the weather?", "role": "user"} + ] assert props["$ai_output_choices"] == [ {"content": "It's sunny!", "role": "assistant", "finish_reason": "stop"} ] @@ -440,7 +452,10 @@ def test_does_not_overwrite_existing_genai_attributes(self, mock_client): transformed = exporter._transform_span(span) # Should preserve existing GenAI attribute, not overwrite - assert transformed.attributes["gen_ai.tool.call.arguments"] == '{"existing": "value"}' + assert ( + transformed.attributes["gen_ai.tool.call.arguments"] + == '{"existing": "value"}' + ) def test_tool_span_export_with_mapped_attributes(self, mock_client): exporter = PydanticAISpanExporter(mock_client, distinct_id="user_123") From f575666af307baacab14f2e992c3b4cae021df8e Mon Sep 17 00:00:00 2001 From: Andrew Maguire Date: Tue, 2 Dec 2025 16:32:00 +0000 Subject: [PATCH 06/14] Update posthog/ai/otel/exporter.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- posthog/ai/otel/exporter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/posthog/ai/otel/exporter.py b/posthog/ai/otel/exporter.py index 8bbdde7c..8ca451f5 100644 --- a/posthog/ai/otel/exporter.py +++ b/posthog/ai/otel/exporter.py @@ -436,7 +436,7 @@ def _create_trace_event( "$ai_span_name": agent_name, "$ai_latency": latency, "$ai_is_error": is_error, - "$ai_framework": "pydantic-ai", + "$ai_framework": "opentelemetry", **self._properties, } From b4bd909aacef1334c112c4b96185de9f2815d1a3 Mon Sep 17 00:00:00 2001 From: Andrew Maguire Date: Tue, 2 Dec 2025 16:33:55 +0000 Subject: [PATCH 07/14] fix: Export PydanticAISpanExporter from package root --- posthog/ai/pydantic_ai/__init__.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/posthog/ai/pydantic_ai/__init__.py b/posthog/ai/pydantic_ai/__init__.py index e670c503..ed9973d8 100644 --- a/posthog/ai/pydantic_ai/__init__.py +++ b/posthog/ai/pydantic_ai/__init__.py @@ -5,6 +5,7 @@ with PostHog tracing. """ +from posthog.ai.pydantic_ai.exporter import PydanticAISpanExporter from posthog.ai.pydantic_ai.instrument import instrument_pydantic_ai -__all__ = ["instrument_pydantic_ai"] +__all__ = ["instrument_pydantic_ai", "PydanticAISpanExporter"] From e433db047b6ed7c84907f713c4fd7496eaee3108 Mon Sep 17 00:00:00 2001 From: Andrew Maguire Date: Tue, 2 Dec 2025 16:40:23 +0000 Subject: [PATCH 08/14] fix: Emit $ai_trace events for agent spans instead of skipping Agent spans now create $ai_trace events so that child spans (tools, etc.) have a valid parent reference. Previously agent spans were skipped, which caused orphaned $ai_parent_id references in tool spans. --- posthog/ai/otel/exporter.py | 13 ++++++++++--- posthog/test/ai/otel/test_exporter.py | 15 ++++++++++----- 2 files changed, 20 insertions(+), 8 deletions(-) diff --git a/posthog/ai/otel/exporter.py b/posthog/ai/otel/exporter.py index 8ca451f5..49b373b8 100644 --- a/posthog/ai/otel/exporter.py +++ b/posthog/ai/otel/exporter.py @@ -231,10 +231,17 @@ def _span_to_event(self, span: ReadableSpan) -> Optional[Dict[str, Any]]: error_message, ) - # Agent run span → skip (PostHog UI auto-creates trace wrapper from generation events) - # The $ai_trace_id on generation events is sufficient for grouping + # Agent run span → $ai_trace if self._is_agent_span(span_name, attrs): - return None # Don't emit separate $ai_trace events + return self._create_trace_event( + span, + attrs, + trace_id, + span_id, + latency, + is_error, + error_message, + ) # Tool execution span → $ai_span if self._is_tool_span(span_name, attrs): diff --git a/posthog/test/ai/otel/test_exporter.py b/posthog/test/ai/otel/test_exporter.py index 819dc291..c280f456 100644 --- a/posthog/test/ai/otel/test_exporter.py +++ b/posthog/test/ai/otel/test_exporter.py @@ -234,9 +234,9 @@ def test_generation_event_with_model_parameters(self, mock_client): class TestAgentSpanHandling: - """Tests for agent span handling (should be skipped).""" + """Tests for agent span handling ($ai_trace events).""" - def test_agent_span_is_skipped(self, mock_client): + def test_agent_span_creates_trace_event(self, mock_client): exporter = PostHogSpanExporter(mock_client, distinct_id="user_123") span = create_mock_span( @@ -245,16 +245,21 @@ def test_agent_span_is_skipped(self, mock_client): exporter.export([span]) - mock_client.capture.assert_not_called() + mock_client.capture.assert_called_once() + call_kwargs = mock_client.capture.call_args[1] + assert call_kwargs["event"] == "$ai_trace" + assert call_kwargs["properties"]["$ai_span_name"] == "TestAgent" - def test_invoke_agent_span_is_skipped(self, mock_client): + def test_invoke_agent_span_creates_trace_event(self, mock_client): exporter = PostHogSpanExporter(mock_client, distinct_id="user_123") span = create_mock_span(name="invoke_agent", attributes={}) exporter.export([span]) - mock_client.capture.assert_not_called() + mock_client.capture.assert_called_once() + call_kwargs = mock_client.capture.call_args[1] + assert call_kwargs["event"] == "$ai_trace" class TestToolSpanEventCreation: From eee31fc9daa56d05169b4a40a0ce3f41cf4602d9 Mon Sep 17 00:00:00 2001 From: Andrew Maguire Date: Tue, 2 Dec 2025 16:43:33 +0000 Subject: [PATCH 09/14] fix: inherit privacy_mode from client when not explicitly set --- posthog/ai/pydantic_ai/instrument.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/posthog/ai/pydantic_ai/instrument.py b/posthog/ai/pydantic_ai/instrument.py index 7aa254cf..0492edf5 100644 --- a/posthog/ai/pydantic_ai/instrument.py +++ b/posthog/ai/pydantic_ai/instrument.py @@ -12,7 +12,7 @@ def instrument_pydantic_ai( client: PostHogClient, distinct_id: Optional[str] = None, - privacy_mode: bool = False, + privacy_mode: Optional[bool] = None, properties: Optional[Dict[str, Any]] = None, groups: Optional[Dict[str, Any]] = None, debug: bool = False, @@ -40,7 +40,7 @@ def instrument_pydantic_ai( distinct_id: Default distinct ID for all events. If not provided, events will use the trace ID as distinct_id. privacy_mode: If True, message content will be redacted from events. - Useful for sensitive data. + If not specified, inherits from client.privacy_mode. properties: Additional properties to include in all events groups: PostHog groups to associate with all events debug: Enable debug logging for troubleshooting @@ -68,6 +68,10 @@ def instrument_pydantic_ai( from posthog.ai.pydantic_ai.exporter import PydanticAISpanExporter + # Resolve privacy_mode from client if not explicitly set + if privacy_mode is None: + privacy_mode = getattr(client, "privacy_mode", False) + # Create the Pydantic AI-specific exporter (handles message format normalization) exporter = PydanticAISpanExporter( client=client, From e0be487fa3c82425123ffadbdde72eddb3e43d5a Mon Sep 17 00:00:00 2001 From: Andrew Maguire Date: Tue, 2 Dec 2025 16:47:26 +0000 Subject: [PATCH 10/14] refactor: remove privacy_mode param, always use client setting --- posthog/ai/pydantic_ai/instrument.py | 7 +------ posthog/test/ai/pydantic_ai/test_instrument.py | 9 ++++++--- 2 files changed, 7 insertions(+), 9 deletions(-) diff --git a/posthog/ai/pydantic_ai/instrument.py b/posthog/ai/pydantic_ai/instrument.py index 0492edf5..4ff0aff9 100644 --- a/posthog/ai/pydantic_ai/instrument.py +++ b/posthog/ai/pydantic_ai/instrument.py @@ -12,7 +12,6 @@ def instrument_pydantic_ai( client: PostHogClient, distinct_id: Optional[str] = None, - privacy_mode: Optional[bool] = None, properties: Optional[Dict[str, Any]] = None, groups: Optional[Dict[str, Any]] = None, debug: bool = False, @@ -39,8 +38,6 @@ def instrument_pydantic_ai( client: PostHog client instance for sending events distinct_id: Default distinct ID for all events. If not provided, events will use the trace ID as distinct_id. - privacy_mode: If True, message content will be redacted from events. - If not specified, inherits from client.privacy_mode. properties: Additional properties to include in all events groups: PostHog groups to associate with all events debug: Enable debug logging for troubleshooting @@ -68,9 +65,7 @@ def instrument_pydantic_ai( from posthog.ai.pydantic_ai.exporter import PydanticAISpanExporter - # Resolve privacy_mode from client if not explicitly set - if privacy_mode is None: - privacy_mode = getattr(client, "privacy_mode", False) + privacy_mode = getattr(client, "privacy_mode", False) # Create the Pydantic AI-specific exporter (handles message format normalization) exporter = PydanticAISpanExporter( diff --git a/posthog/test/ai/pydantic_ai/test_instrument.py b/posthog/test/ai/pydantic_ai/test_instrument.py index 8b164187..d80792b5 100644 --- a/posthog/test/ai/pydantic_ai/test_instrument.py +++ b/posthog/test/ai/pydantic_ai/test_instrument.py @@ -41,11 +41,14 @@ def test_basic_instrumentation(self, mock_client): settings = mock_instrument_all.call_args[0][0] assert isinstance(settings, InstrumentationSettings) - def test_privacy_mode_disables_content(self, mock_client): + def test_privacy_mode_disables_content(self): from posthog.ai.pydantic_ai import instrument_pydantic_ai + client = MagicMock() + client.privacy_mode = True + with patch.object(Agent, "instrument_all") as mock_instrument_all: - instrument_pydantic_ai(mock_client, privacy_mode=True) + instrument_pydantic_ai(client) settings = mock_instrument_all.call_args[0][0] assert settings.include_content is False @@ -54,7 +57,7 @@ def test_privacy_mode_false_includes_content(self, mock_client): from posthog.ai.pydantic_ai import instrument_pydantic_ai with patch.object(Agent, "instrument_all") as mock_instrument_all: - instrument_pydantic_ai(mock_client, privacy_mode=False) + instrument_pydantic_ai(mock_client) settings = mock_instrument_all.call_args[0][0] assert settings.include_content is True From ecdbbc8e531320a59755d630b65aec12ec3aac5d Mon Sep 17 00:00:00 2001 From: Andrew Maguire Date: Tue, 2 Dec 2025 16:48:37 +0000 Subject: [PATCH 11/14] docs: update README to reflect privacy_mode from client --- posthog/ai/pydantic_ai/README.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/posthog/ai/pydantic_ai/README.md b/posthog/ai/pydantic_ai/README.md index edb3fc06..a4876648 100644 --- a/posthog/ai/pydantic_ai/README.md +++ b/posthog/ai/pydantic_ai/README.md @@ -40,7 +40,6 @@ result = await agent.run("What's the weather in San Francisco?") instrument_pydantic_ai( client=posthog, # PostHog client instance distinct_id="user_123", # User identifier for events - privacy_mode=False, # Set True to exclude message content properties={ # Additional properties for all events "$ai_session_id": "session_abc", }, @@ -51,13 +50,15 @@ instrument_pydantic_ai( ) ``` +Privacy mode is inherited from the client - set `privacy_mode=True` when creating your PostHog client to exclude message content. + ## What Gets Captured ### Model Calls (`$ai_generation` events) Every LLM API call creates an event with: - Model name and provider -- Input/output messages (unless `privacy_mode=True`) +- Input/output messages (unless privacy mode is enabled on the client) - Token usage (input, output) - Latency - Error status From 0a77c1527f1404963c39dab47690cb98e1d4cab4 Mon Sep 17 00:00:00 2001 From: Andrew Maguire Date: Tue, 2 Dec 2025 16:50:51 +0000 Subject: [PATCH 12/14] fix: remove unused imports in test files --- posthog/test/ai/otel/test_exporter.py | 4 ++-- posthog/test/ai/pydantic_ai/test_exporter.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/posthog/test/ai/otel/test_exporter.py b/posthog/test/ai/otel/test_exporter.py index c280f456..9db96ffa 100644 --- a/posthog/test/ai/otel/test_exporter.py +++ b/posthog/test/ai/otel/test_exporter.py @@ -3,14 +3,14 @@ """ import json -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock import pytest try: from opentelemetry.sdk.trace import ReadableSpan from opentelemetry.sdk.trace.export import SpanExportResult - from opentelemetry.trace import SpanContext, Status, StatusCode, TraceFlags + from opentelemetry.trace import StatusCode from posthog.ai.otel import PostHogSpanExporter diff --git a/posthog/test/ai/pydantic_ai/test_exporter.py b/posthog/test/ai/pydantic_ai/test_exporter.py index 6c7f5584..82ae13e1 100644 --- a/posthog/test/ai/pydantic_ai/test_exporter.py +++ b/posthog/test/ai/pydantic_ai/test_exporter.py @@ -3,14 +3,14 @@ """ import json -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock import pytest try: from opentelemetry.sdk.trace import ReadableSpan from opentelemetry.sdk.trace.export import SpanExportResult - from opentelemetry.trace import SpanContext, Status, StatusCode, TraceFlags + from opentelemetry.trace import StatusCode from posthog.ai.pydantic_ai.exporter import PydanticAISpanExporter From 8f0249add857a6c29312fd6e6777d959a8d890bc Mon Sep 17 00:00:00 2001 From: Andrew Maguire Date: Tue, 2 Dec 2025 17:20:22 +0000 Subject: [PATCH 13/14] fix: resolve mypy errors in otel and pydantic_ai exporters --- posthog/ai/otel/exporter.py | 18 ++++++++++-------- posthog/ai/pydantic_ai/exporter.py | 13 ++++++++----- 2 files changed, 18 insertions(+), 13 deletions(-) diff --git a/posthog/ai/otel/exporter.py b/posthog/ai/otel/exporter.py index 49b373b8..d018e080 100644 --- a/posthog/ai/otel/exporter.py +++ b/posthog/ai/otel/exporter.py @@ -8,23 +8,25 @@ import json import logging -from typing import Any, Dict, Optional, Sequence, Union +from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence, Union from posthog.client import Client as PostHogClient +if TYPE_CHECKING: + from opentelemetry.sdk.trace import ReadableSpan + from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult + from opentelemetry.trace import StatusCode + try: from opentelemetry.sdk.trace import ReadableSpan from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult from opentelemetry.trace import StatusCode OTEL_AVAILABLE = True + _BASE_CLASS = SpanExporter except ImportError: OTEL_AVAILABLE = False - # Define stub types for type hints when OTel is not installed - ReadableSpan = Any # type: ignore - SpanExporter = object # type: ignore - SpanExportResult = Any # type: ignore - StatusCode = Any # type: ignore + _BASE_CLASS = object logger = logging.getLogger(__name__) @@ -77,7 +79,7 @@ class GenAIAttributes: SERVER_PORT = "server.port" -class PostHogSpanExporter(SpanExporter if OTEL_AVAILABLE else object): +class PostHogSpanExporter(_BASE_CLASS): # type: ignore[valid-type,misc] """ OpenTelemetry SpanExporter that sends AI/LLM spans to PostHog. @@ -284,7 +286,7 @@ def _is_agent_span(self, span_name: str, attrs: Dict[str, Any]) -> bool: """Check if span represents an agent run.""" return span_name in ("agent run", "invoke_agent") or attrs.get( GenAIAttributes.AGENT_NAME - ) + ) is not None def _is_tool_span(self, span_name: str, attrs: Dict[str, Any]) -> bool: """Check if span represents a tool/function execution.""" diff --git a/posthog/ai/pydantic_ai/exporter.py b/posthog/ai/pydantic_ai/exporter.py index a5e56bdd..9564c7d2 100644 --- a/posthog/ai/pydantic_ai/exporter.py +++ b/posthog/ai/pydantic_ai/exporter.py @@ -5,24 +5,27 @@ Pydantic AI-specific transformations like message format normalization. """ -from typing import Any, Dict, List, Optional, Sequence +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence + +if TYPE_CHECKING: + from opentelemetry.sdk.trace import ReadableSpan + from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult try: from opentelemetry.sdk.trace import ReadableSpan from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult OTEL_AVAILABLE = True + _BASE_CLASS = SpanExporter except ImportError: OTEL_AVAILABLE = False - ReadableSpan = Any # type: ignore - SpanExporter = object # type: ignore - SpanExportResult = Any # type: ignore + _BASE_CLASS = object from posthog.ai.otel import PostHogSpanExporter from posthog.client import Client as PostHogClient -class PydanticAISpanExporter(SpanExporter if OTEL_AVAILABLE else object): +class PydanticAISpanExporter(_BASE_CLASS): # type: ignore[valid-type,misc] """ SpanExporter for Pydantic AI that normalizes messages to OpenAI format. From dc7af941907ec6b267a722ae139ac17cf1bc1568 Mon Sep 17 00:00:00 2001 From: Andrew Maguire Date: Tue, 2 Dec 2025 17:28:04 +0000 Subject: [PATCH 14/14] style: format exporter.py --- posthog/ai/otel/exporter.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/posthog/ai/otel/exporter.py b/posthog/ai/otel/exporter.py index d018e080..ff529418 100644 --- a/posthog/ai/otel/exporter.py +++ b/posthog/ai/otel/exporter.py @@ -284,9 +284,10 @@ def _is_generation_span(self, span_name: str, attrs: Dict[str, Any]) -> bool: def _is_agent_span(self, span_name: str, attrs: Dict[str, Any]) -> bool: """Check if span represents an agent run.""" - return span_name in ("agent run", "invoke_agent") or attrs.get( - GenAIAttributes.AGENT_NAME - ) is not None + return ( + span_name in ("agent run", "invoke_agent") + or attrs.get(GenAIAttributes.AGENT_NAME) is not None + ) def _is_tool_span(self, span_name: str, attrs: Dict[str, Any]) -> bool: """Check if span represents a tool/function execution."""