style: format exporter.py

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/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`.

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"]

posthog/ai/otel/exporter.py

@@ -0,0 +1,559 @@
+"""
+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 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
+    _BASE_CLASS = object
+
+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(_BASE_CLASS):  # type: ignore[valid-type,misc]
+    """
+    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 → $ai_trace
+        if self._is_agent_span(span_name, attrs):
+            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):
+            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) is not None
+        )
+
+    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": "opentelemetry",
+            **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": "opentelemetry",
+            **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": "opentelemetry",
+            **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": "opentelemetry",
+            **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:]}"

posthog/ai/pydantic_ai/README.md

@@ -0,0 +1,135 @@
+# 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
+    properties={              # Additional properties for all events
+        "$ai_session_id": "session_abc",
+    },
+    groups={                  # PostHog groups
+        "company": "acme",
+    },
+    debug=False,              # Enable debug logging
+)
+```
+
+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 is enabled on the client)
+- 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,
+))
+```

posthog/ai/pydantic_ai/__init__.py

@@ -0,0 +1,11 @@
+"""
+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.exporter import PydanticAISpanExporter
+from posthog.ai.pydantic_ai.instrument import instrument_pydantic_ai
+
+__all__ = ["instrument_pydantic_ai", "PydanticAISpanExporter"]

posthog/ai/pydantic_ai/exporter.py

@@ -0,0 +1,270 @@
+"""
+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 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
+    _BASE_CLASS = object
+
+from posthog.ai.otel import PostHogSpanExporter
+from posthog.client import Client as PostHogClient
+
+
+class PydanticAISpanExporter(_BASE_CLASS):  # type: ignore[valid-type,misc]
+    """
+    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

posthog/ai/pydantic_ai/instrument.py

@@ -0,0 +1,91 @@
+"""
+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,
+    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.
+        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
+
+    privacy_mode = getattr(client, "privacy_mode", False)
+
+    # 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)

posthog/test/ai/otel/test_exporter.py

@@ -0,0 +1,484 @@
+"""
+Tests for PostHogSpanExporter - the generic OpenTelemetry span exporter for PostHog.
+"""
+
+import json
+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 StatusCode
+
+    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 ($ai_trace events)."""
+
+    def test_agent_span_creates_trace_event(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_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_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_called_once()
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["event"] == "$ai_trace"
+
+
+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

posthog/test/ai/pydantic_ai/test_exporter.py

@@ -0,0 +1,552 @@
+"""
+Tests for PydanticAISpanExporter - handles Pydantic AI message format normalization.
+"""
+
+import json
+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 StatusCode
+
+    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

posthog/test/ai/pydantic_ai/test_instrument.py

@@ -0,0 +1,100 @@
+"""
+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):
+        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(client)
+
+            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)
+
+            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()

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",
 ]