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

.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,24 @@
+# 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
+
+## Run
+
+```bash
+source .env
+python simple_query.py
+python instrument_reuse.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/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,125 @@
+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.processor import PostHogClaudeAgentProcessor
+
+__all__ = ["PostHogClaudeAgentProcessor", "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