feat(ai): add Claude Agent SDK integration for LLM analytics (#477)

* feat(ai): add Claude Agent SDK integration for LLM analytics

Add posthog.ai.claude_agent_sdk module that wraps claude_agent_sdk.query()
to automatically emit $ai_generation, $ai_span, and $ai_trace events.

- PostHogClaudeAgentProcessor with _GenerationTracker that reconstructs
  per-turn generation metrics from Anthropic StreamEvents
- Two entry points: query() drop-in replacement and instrument() for
  configure-once reuse
- Two-slot input tracking to correctly associate tool results with
  subsequent generations despite SDK message ordering
- All instrumentation wrapped in try/except so PostHog errors never
  interrupt the underlying Claude Agent SDK query
- 16 unit tests covering generation, multi-turn, fallback, tool spans,
  traces, privacy mode, personless mode, custom properties
- Example scripts (simple_query.py, instrument_reuse.py)

* fix(ai): address review feedback on claude agent sdk integration

- Honor per-call privacy override for $ai_input/$ai_output_choices
  in generation events (was only checking instance-level privacy)
- Pass groups directly to _capture_event instead of fragile
  save/restore pattern on self._groups (thread-safe, exception-safe)
- Fix tool span parent linkage: use tracker.current_span_id for
  in-progress generation instead of stale current_generation_span_id

* style: ruff format claude_agent_sdk files

* feat(ai): add PostHogClaudeSDKClient for stateful multi-turn conversations

Wraps ClaudeSDKClient to instrument receive_response() with the same
generation/span/trace tracking as query(). Supports multi-turn
conversations with full history — each turn emits its own $ai_generation
events, all linked by a shared $ai_trace_id. The $ai_trace event is
emitted on disconnect() to cover the entire session.

Usage:
    async with PostHogClaudeSDKClient(options, posthog_client=ph) as client:
        await client.query("Hello")
        async for msg in client.receive_response():
            ...
        await client.query("Follow up")  # has conversation history
        async for msg in client.receive_response():
            ...

* docs(ai): add multi-turn conversation example for Claude Agent SDK

* refactor(ai): split PostHogClaudeSDKClient into client.py

* style: ruff format multi_turn.py example

* fix: resolve ruff lint errors in claude agent sdk files

Author: andrewm4894

.sampo/changesets/claude-agent-sdk-integration.md

@@ -0,0 +1,5 @@
+---
+pypi/posthog: minor
+---
+
+feat(ai): add Claude Agent SDK integration for LLM analytics

examples/example-ai-claude-agent-sdk/.env.example

@@ -0,0 +1,3 @@
+POSTHOG_API_KEY=phc_your_project_api_key
+POSTHOG_HOST=https://us.i.posthog.com
+ANTHROPIC_API_KEY=sk-ant-your_api_key

examples/example-ai-claude-agent-sdk/README.md

@@ -0,0 +1,26 @@
+# Claude Agent SDK + PostHog AI Examples
+
+Track Claude Agent SDK calls with PostHog.
+
+## Setup
+
+```bash
+pip install -r requirements.txt
+cp .env.example .env
+# Fill in your API keys in .env
+```
+
+## Examples
+
+- **simple_query.py** - Single query using the `query()` drop-in replacement
+- **instrument_reuse.py** - Configure-once with `instrument()`, reuse across multiple queries
+- **multi_turn.py** - Multi-turn conversation with history using `PostHogClaudeSDKClient`
+
+## Run
+
+```bash
+source .env
+python simple_query.py
+python instrument_reuse.py
+python multi_turn.py
+```

examples/example-ai-claude-agent-sdk/instrument_reuse.py

@@ -0,0 +1,43 @@
+"""Claude Agent SDK with instrument() for reusable config, tracked by PostHog."""
+
+import asyncio
+import os
+
+from claude_agent_sdk import ClaudeAgentOptions, AssistantMessage, TextBlock
+from posthog import Posthog
+from posthog.ai.claude_agent_sdk import instrument
+
+posthog = Posthog(
+    os.environ["POSTHOG_API_KEY"],
+    host=os.environ.get("POSTHOG_HOST", "https://us.i.posthog.com"),
+)
+
+# Configure once, reuse for multiple queries
+ph = instrument(
+    client=posthog,
+    distinct_id="example-user",
+    properties={"app": "demo", "environment": "development"},
+)
+
+
+async def ask(prompt: str) -> None:
+    print(f"\n> {prompt}")
+    options = ClaudeAgentOptions(
+        max_turns=2,
+        permission_mode="plan",
+    )
+
+    async for message in ph.query(prompt=prompt, options=options):
+        if isinstance(message, AssistantMessage):
+            for block in message.content:
+                if isinstance(block, TextBlock):
+                    print(f"  {block.text}")
+
+
+async def main():
+    await ask("What is the capital of France? Reply in one sentence.")
+    await ask("What is 15% of 280? Reply in one sentence.")
+
+
+asyncio.run(main())
+posthog.shutdown()

examples/example-ai-claude-agent-sdk/multi_turn.py

@@ -0,0 +1,68 @@
+"""Claude Agent SDK multi-turn conversation with history, tracked by PostHog."""
+
+import asyncio
+import os
+
+from claude_agent_sdk import AssistantMessage, ResultMessage
+from claude_agent_sdk.types import ClaudeAgentOptions, TextBlock
+from posthog import Posthog
+from posthog.ai.claude_agent_sdk import PostHogClaudeSDKClient
+
+posthog = Posthog(
+    os.environ["POSTHOG_API_KEY"],
+    host=os.environ.get("POSTHOG_HOST", "https://us.i.posthog.com"),
+)
+
+
+async def main():
+    options = ClaudeAgentOptions(
+        max_turns=5,
+        permission_mode="plan",
+    )
+
+    async with PostHogClaudeSDKClient(
+        options,
+        posthog_client=posthog,
+        posthog_distinct_id="example-user",
+        posthog_properties={"example": "multi_turn"},
+    ) as client:
+        # Turn 1
+        print("> What is the capital of France?")
+        await client.query("What is the capital of France? Reply in one sentence.")
+        async for message in client.receive_response():
+            if isinstance(message, AssistantMessage):
+                for block in message.content:
+                    if isinstance(block, TextBlock):
+                        print(f"  {block.text}")
+            elif isinstance(message, ResultMessage):
+                print(f"  [{message.num_turns} turns, ${message.total_cost_usd:.4f}]")
+
+        # Turn 2 — has full conversation history
+        print("\n> And what language do they speak there?")
+        await client.query(
+            "And what language do they speak there? Reply in one sentence."
+        )
+        async for message in client.receive_response():
+            if isinstance(message, AssistantMessage):
+                for block in message.content:
+                    if isinstance(block, TextBlock):
+                        print(f"  {block.text}")
+            elif isinstance(message, ResultMessage):
+                print(f"  [{message.num_turns} turns, ${message.total_cost_usd:.4f}]")
+
+        # Turn 3 — still has context from both previous turns
+        print("\n> How do you say 'hello' in that language?")
+        await client.query(
+            "How do you say 'hello' in that language? Reply in one sentence."
+        )
+        async for message in client.receive_response():
+            if isinstance(message, AssistantMessage):
+                for block in message.content:
+                    if isinstance(block, TextBlock):
+                        print(f"  {block.text}")
+            elif isinstance(message, ResultMessage):
+                print(f"  [{message.num_turns} turns, ${message.total_cost_usd:.4f}]")
+
+
+asyncio.run(main())
+posthog.shutdown()

examples/example-ai-claude-agent-sdk/requirements.txt

@@ -0,0 +1,2 @@
+posthog>=7.9.12
+claude-agent-sdk

examples/example-ai-claude-agent-sdk/simple_query.py

@@ -0,0 +1,33 @@
+"""Claude Agent SDK simple query, tracked by PostHog."""
+
+import asyncio
+import os
+
+from claude_agent_sdk import ClaudeAgentOptions
+from posthog import Posthog
+from posthog.ai.claude_agent_sdk import query
+
+posthog = Posthog(
+    os.environ["POSTHOG_API_KEY"],
+    host=os.environ.get("POSTHOG_HOST", "https://us.i.posthog.com"),
+)
+
+
+async def main():
+    options = ClaudeAgentOptions(
+        max_turns=2,
+        permission_mode="plan",
+    )
+
+    async for message in query(
+        prompt="What is 2 + 2? Reply in one sentence.",
+        options=options,
+        posthog_client=posthog,
+        posthog_distinct_id="example-user",
+        posthog_properties={"example": "simple_query"},
+    ):
+        print(f"[{type(message).__name__}]")
+
+
+asyncio.run(main())
+posthog.shutdown()

examples/example-ai-claude-agent-sdk/uv.toml

@@ -0,0 +1 @@
+exclude-newer = "7 days"

posthog/ai/claude_agent_sdk/__init__.py

@@ -0,0 +1,133 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union
+
+if TYPE_CHECKING:
+    from claude_agent_sdk.types import ClaudeAgentOptions, ResultMessage
+
+    from posthog.client import Client
+
+try:
+    import claude_agent_sdk  # noqa: F401
+except ImportError:
+    raise ModuleNotFoundError(
+        "Please install the Claude Agent SDK to use this feature: 'pip install claude-agent-sdk'"
+    )
+
+from posthog.ai.claude_agent_sdk.client import PostHogClaudeSDKClient
+from posthog.ai.claude_agent_sdk.processor import PostHogClaudeAgentProcessor
+
+__all__ = [
+    "PostHogClaudeAgentProcessor",
+    "PostHogClaudeSDKClient",
+    "instrument",
+    "query",
+]
+
+
+def instrument(
+    client: Optional[Client] = None,
+    distinct_id: Optional[Union[str, Callable[[ResultMessage], Optional[str]]]] = None,
+    privacy_mode: bool = False,
+    groups: Optional[Dict[str, Any]] = None,
+    properties: Optional[Dict[str, Any]] = None,
+) -> PostHogClaudeAgentProcessor:
+    """
+    Create a PostHog-instrumented query wrapper for the Claude Agent SDK.
+
+    Returns a PostHogClaudeAgentProcessor whose .query() method is a drop-in
+    replacement for claude_agent_sdk.query() that automatically emits
+    $ai_generation, $ai_span, and $ai_trace events.
+
+    Args:
+        client: Optional PostHog client instance. If not provided, uses the default client.
+        distinct_id: Optional distinct ID to associate with all events.
+            Can also be a callable that takes a ResultMessage and returns a distinct ID.
+        privacy_mode: If True, redacts sensitive information in tracking.
+        groups: Optional PostHog groups to associate with events.
+        properties: Optional additional properties to include with all events.
+
+    Returns:
+        PostHogClaudeAgentProcessor: A processor whose .query() method wraps claude_agent_sdk.query().
+
+    Example:
+        ```python
+        from posthog.ai.claude_agent_sdk import instrument
+
+        ph = instrument(distinct_id="my-app", properties={"env": "prod"})
+
+        async for message in ph.query(prompt="Hello", options=options):
+            print(message)
+        ```
+    """
+    return PostHogClaudeAgentProcessor(
+        client=client,
+        distinct_id=distinct_id,
+        privacy_mode=privacy_mode,
+        groups=groups,
+        properties=properties,
+    )
+
+
+async def query(
+    *,
+    prompt: Any,
+    options: Optional[ClaudeAgentOptions] = None,
+    transport: Any = None,
+    posthog_client: Optional[Client] = None,
+    posthog_distinct_id: Optional[
+        Union[str, Callable[[ResultMessage], Optional[str]]]
+    ] = None,
+    posthog_trace_id: Optional[str] = None,
+    posthog_properties: Optional[Dict[str, Any]] = None,
+    posthog_privacy_mode: bool = False,
+    posthog_groups: Optional[Dict[str, Any]] = None,
+):
+    """
+    Drop-in replacement for claude_agent_sdk.query() with PostHog instrumentation.
+
+    All original messages are yielded unchanged. PostHog events ($ai_generation,
+    $ai_span, $ai_trace) are emitted automatically.
+
+    Args:
+        prompt: The prompt (same as claude_agent_sdk.query)
+        options: ClaudeAgentOptions (same as claude_agent_sdk.query)
+        transport: Optional transport (same as claude_agent_sdk.query)
+        posthog_client: Optional PostHog client instance.
+        posthog_distinct_id: Optional distinct ID for this query.
+        posthog_trace_id: Optional trace ID (auto-generated if not provided).
+        posthog_properties: Extra properties to include with all events.
+        posthog_privacy_mode: If True, redacts sensitive content.
+        posthog_groups: Optional PostHog groups.
+
+    Example:
+        ```python
+        from posthog.ai.claude_agent_sdk import query
+
+        async for message in query(
+            prompt="Hello",
+            options=options,
+            posthog_distinct_id="my-app",
+            posthog_properties={"pr_number": 123},
+        ):
+            print(message)
+        ```
+    """
+    processor = PostHogClaudeAgentProcessor(
+        client=posthog_client,
+        distinct_id=posthog_distinct_id,
+        privacy_mode=posthog_privacy_mode,
+        groups=posthog_groups,
+        properties={},
+    )
+
+    async for message in processor.query(
+        prompt=prompt,
+        options=options,
+        transport=transport,
+        posthog_trace_id=posthog_trace_id,
+        posthog_properties=posthog_properties,
+        posthog_privacy_mode=posthog_privacy_mode,
+        posthog_groups=posthog_groups,
+    ):
+        yield message