WIP

Restores the process_exception method that was removed in v6.7.5 (PR #328),
which broke exception capture from Django views and downstream middleware.

Django converts view exceptions into responses before they propagate through
the middleware stack's __call__ method, so the context manager's exception
handler never sees them. Django provides these exceptions via the
process_exception hook instead.

Changes:
- Add process_exception method to capture exceptions from views and downstream
  middleware with proper request context and tags
- Add tests verifying process_exception behavior and settings (capture_exceptions,
  request_filter)

.github/dependabot.yml

@@ -0,0 +1,36 @@
+version: 2
+updates:
+  - package-ecosystem: "pip"
+    directory: "/"
+    schedule:
+      interval: "daily"
+      time: "10:00"
+      timezone: "UTC"
+    groups:
+      ai-providers:
+        patterns:
+          - "openai"
+          - "anthropic"
+          - "google-genai"
+          - "langchain-core"
+          - "langchain-community"
+          - "langchain-openai"
+          - "langchain-anthropic"
+          - "langgraph"
+    allow:
+      - dependency-name: "openai"
+      - dependency-name: "anthropic"
+      - dependency-name: "google-genai"
+      - dependency-name: "langchain-core"
+      - dependency-name: "langchain-community"
+      - dependency-name: "langchain-openai"
+      - dependency-name: "langchain-anthropic"
+      - dependency-name: "langgraph"
+    open-pull-requests-limit: 1
+    reviewers:
+      - "PostHog/team-llm-analytics"
+    # Uncomment below to enable auto-merge for minor updates when CI passes
+    # pull-request-branch-name:
+    #   separator: "/"
+    # assignees:
+    #   - "PostHog/ai-team"

.github/workflows/generate-references.yml

@@ -0,0 +1,48 @@
+name: "Generate References"
+
+on:
+  workflow_dispatch:
+
+jobs:
+  docs-generation:
+    name: Generate references
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout the repository
+        uses: actions/checkout@85e6279cec87321a52edac9c87bce653a07cf6c2
+        with:
+          fetch-depth: 0
+          token: ${{ secrets.POSTHOG_BOT_PAT }}
+
+      - name: Set up Python
+        uses: actions/setup-python@8d9ed9ac5c53483de85588cdf95a591a75ab9f55
+        with:
+          python-version: 3.11.11
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@0c5e2b8115b80b4c7c5ddf6ffdd634974642d182 # v5.4.1
+        with:
+            enable-cache: true
+            pyproject-file: 'pyproject.toml'
+
+      - name: Generate references
+        run: |
+          uv run bin/docs generate-references
+
+      - name: Check for changes in references
+        id: changes
+        run: |
+          if [ -n "$(git status --porcelain references/)" ]; then
+            echo "changed=true" >> $GITHUB_OUTPUT
+            echo "New references generated in references directory:"
+            git status --porcelain references/
+          else
+            echo "changed=false" >> $GITHUB_OUTPUT
+            echo "No new references generated in references directory"
+          fi
+
+      - uses: stefanzweifel/git-auto-commit-action@778341af668090896ca464160c2def5d1d1a3eb0
+        if: steps.changes.outputs.changed == 'true'
+        with:
+          commit_message: "Update generated references"
+          file_pattern: references/

.github/workflows/release.yml

@@ -20,7 +20,7 @@ jobs:
         uses: actions/checkout@85e6279cec87321a52edac9c87bce653a07cf6c2
         with:
           fetch-depth: 0
-          token: ${{ secrets.POSTHOG_BOT_GITHUB_TOKEN }}
+          token: ${{ secrets.POSTHOG_BOT_PAT }}
 
       - name: Set up Python
         uses: actions/setup-python@8d9ed9ac5c53483de85588cdf95a591a75ab9f55
@@ -45,7 +45,13 @@ jobs:
       - name: Create GitHub release
         uses: actions/create-release@0cb9c9b65d5d1901c1f53e5e66eaf4afd303e70e # v1
         env:
-          GITHUB_TOKEN: ${{ secrets.POSTHOG_BOT_GITHUB_TOKEN }}
+          GITHUB_TOKEN: ${{ secrets.POSTHOG_BOT_PAT }}
         with:
           tag_name: v${{ env.REPO_VERSION }}
           release_name: ${{ env.REPO_VERSION }}
+      
+      - name: Dispatch generate-references for posthog-python
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        run: |
+          gh workflow run generate-references.yml --ref master

CHANGELOG.md

@@ -1,3 +1,52 @@
+# Unreleased
+
+- fix(django): Restore process_exception method to capture view and downstream middleware exceptions (fixes #329)
+
+# 6.7.11 - 2025-10-28
+
+- feat(ai): Add `$ai_framework` property for framework integrations (e.g. LangChain)
+
+# 6.7.10 - 2025-10-24
+
+- fix(django): Make middleware truly hybrid - compatible with both sync (WSGI) and async (ASGI) Django stacks without breaking sync-only deployments
+
+# 6.7.9 - 2025-10-22
+
+- fix(flags): multi-condition flags with static cohorts returning wrong variants
+
+# 6.7.8 - 2025-10-16
+
+- fix(llma): missing async for OpenAI's streaming implementation
+
+# 6.7.7 - 2025-10-14
+
+- fix: remove deprecated attribute $exception_personURL from exception events
+
+# 6.7.6 - 2025-09-16
+
+- fix: don't sort condition sets with variant overrides to the top
+- fix: Prevent core Client methods from raising exceptions
+
+# 6.7.5 - 2025-09-16
+
+- feat: Django middleware now supports async request handling.
+
+# 6.7.4 - 2025-09-05
+
+- fix: Missing system prompts for some providers
+
+# 6.7.3 - 2025-09-04
+
+- fix: missing usage tokens in Gemini
+
+# 6.7.2 - 2025-09-03
+
+- fix: tool call results in streaming providers
+
+# 6.7.1 - 2025-09-01
+
+- fix: Add base64 inline image sanitization
+
 # 6.7.0 - 2025-08-26
 
 - feat: Add support for feature flag dependencies

bin/docs_scripts/doc_constant.py

@@ -3,6 +3,7 @@ Constants for PostHog Python SDK documentation generation.
 """
 
 from typing import Dict, Union
+from posthog.version import VERSION
 
 # Documentation generation metadata
 DOCUMENTATION_METADATA = {
@@ -27,8 +28,9 @@ DOCSTRING_PATTERNS = {
 
 # Output file configuration
 OUTPUT_CONFIG: Dict[str, Union[str, int]] = {
-    "output_dir": ".",
-    "filename": "posthog-python-references.json",
+    "output_dir": "./references",
+    "filename": f"posthog-python-references-{VERSION}.json",
+    "filename_latest": "posthog-python-references-latest.json",
     "indent": 2,
 }
 

bin/docs_scripts/generate_json_schemas.py

@@ -460,12 +460,23 @@ if __name__ == "__main__":
     try:
         documentation = generate_sdk_documentation()
 
-        # Write to file
+        # Ensure output directory exists
+        output_dir = str(OUTPUT_CONFIG["output_dir"])
+        os.makedirs(output_dir, exist_ok=True)
+
         output_file = os.path.join(
             str(OUTPUT_CONFIG["output_dir"]), str(OUTPUT_CONFIG["filename"])
         )
+        output_file_latest = os.path.join(
+            str(OUTPUT_CONFIG["output_dir"]), str(OUTPUT_CONFIG["filename_latest"])
+        )
+
+        # Write to current version
         with open(output_file, "w") as f:
             json.dump(documentation, f, indent=int(OUTPUT_CONFIG["indent"]))
+        # Write to latest
+        with open(output_file_latest, "w") as f:
+            json.dump(documentation, f, indent=int(OUTPUT_CONFIG["indent"]))
 
         print(f"✓ Generated {output_file}")
 

mypy-baseline.txt

@@ -36,10 +36,5 @@ posthog/client.py:0: error: "None" has no attribute "start"  [attr-defined]
 posthog/client.py:0: error: "None" has no attribute "get"  [attr-defined]
 posthog/client.py:0: error: Statement is unreachable  [unreachable]
 posthog/client.py:0: error: Statement is unreachable  [unreachable]
-posthog/ai/utils.py:0: error: Need type annotation for "output" (hint: "output: list[<type>] = ...")  [var-annotated]
-posthog/ai/utils.py:0: error: Function "builtins.any" is not valid as a type  [valid-type]
-posthog/ai/utils.py:0: note: Perhaps you meant "typing.Any" instead of "any"?
-posthog/ai/utils.py:0: error: Function "builtins.any" is not valid as a type  [valid-type]
-posthog/ai/utils.py:0: note: Perhaps you meant "typing.Any" instead of "any"?
 posthog/client.py:0: error: Name "urlparse" already defined (possibly by an import)  [no-redef]
 posthog/client.py:0: error: Name "parse_qs" already defined (possibly by an import)  [no-redef]

posthog/__init__.py

@@ -2,7 +2,13 @@ import datetime  # noqa: F401
 from typing import Callable, Dict, Optional, Any  # noqa: F401
 from typing_extensions import Unpack
 
-from posthog.args import OptionalCaptureArgs, OptionalSetArgs, ExceptionArg
+from posthog.args import (
+    OptionalCaptureArgs,
+    OptionalSetArgs,
+    OptionalCaptureAIArgs,
+    AI_EVENT_TYPE,
+    ExceptionArg,
+)
 from posthog.client import Client
 from posthog.contexts import (
     new_context as inner_new_context,
@@ -11,7 +17,15 @@ from posthog.contexts import (
     set_context_session as inner_set_context_session,
     identify_context as inner_identify_context,
 )
-from posthog.types import FeatureFlag, FlagsAndPayloads, FeatureFlagResult
+from posthog.feature_flags import (
+    InconclusiveMatchError as InconclusiveMatchError,
+    RequiresServerEvaluation as RequiresServerEvaluation,
+)
+from posthog.types import (
+    FeatureFlag,
+    FlagsAndPayloads,
+    FeatureFlagResult as FeatureFlagResult,
+)
 from posthog.version import VERSION
 
 __version__ = VERSION
@@ -385,6 +399,81 @@ def capture_exception(
     return _proxy("capture_exception", exception=exception, **kwargs)
 
 
+def capture_ai(
+    event: AI_EVENT_TYPE, **kwargs: Unpack[OptionalCaptureAIArgs]
+) -> Optional[str]:
+    """
+    Capture an AI event to the dedicated AI endpoint with support for large payloads.
+
+    This method sends AI events (like $ai_generation, $ai_trace, etc.) to PostHog's
+    specialized AI endpoint (/i/v0/ai) which supports large payloads through multipart/form-data
+    and blob storage in S3.
+
+    Args:
+        event: The AI event type. Must be one of: "$ai_generation", "$ai_trace", "$ai_span",
+               "$ai_embedding", "$ai_metric", "$ai_feedback"
+        distinct_id: The distinct ID of the user.
+        properties: A dictionary of AI event properties. Must include required properties based on event type:
+            - All events: "$ai_model" (required)
+            - $ai_generation: "$ai_provider", "$ai_trace_id" (required)
+            - $ai_trace: "$ai_trace_id" (required)
+            - $ai_span: "$ai_trace_id", "$ai_span_id" (required)
+            - $ai_embedding: "$ai_provider", "$ai_trace_id" (required)
+        blob_properties: List of property names to send as blobs (large data stored in S3).
+            Common blob properties: "$ai_input", "$ai_output_choices", "$ai_input_state", "$ai_output_state"
+            If not provided, defaults to common blob properties based on event type.
+        timestamp: The timestamp of the event.
+        uuid: A unique identifier for the event.
+        groups: A dictionary of group information.
+        disable_geoip: Whether to disable GeoIP for this event.
+
+    Examples:
+        ```python
+        # $ai_generation event with blobs
+        from posthog import capture_ai
+        capture_ai(
+            "$ai_generation",
+            distinct_id="user_123",
+            properties={
+                "$ai_model": "gpt-4",
+                "$ai_provider": "openai",
+                "$ai_trace_id": "trace_abc123",
+                "$ai_input": {
+                    "messages": [
+                        {"role": "user", "content": "Hello!"}
+                    ]
+                },
+                "$ai_output_choices": {
+                    "choices": [{"message": {"role": "assistant", "content": "Hi there!"}}]
+                },
+                "$ai_completion_tokens": 150,
+                "$ai_prompt_tokens": 50
+            },
+            blob_properties=["$ai_input", "$ai_output_choices"]
+        )
+        ```
+
+        ```python
+        # $ai_trace event
+        from posthog import capture_ai
+        capture_ai(
+            "$ai_trace",
+            distinct_id="user_123",
+            properties={
+                "$ai_model": "gpt-4",
+                "$ai_trace_id": "trace_abc123"
+            }
+        )
+        ```
+
+    Category:
+        AI Events
+
+    Note: This method sends events synchronously to the AI endpoint, bypassing the queue system.
+    """
+    return _proxy("capture_ai", event, **kwargs)
+
+
 def feature_enabled(
     key,  # type: str
     distinct_id,  # type: str

posthog/ai/anthropic/__init__.py

@@ -6,6 +6,12 @@ from .anthropic_providers import (
     AsyncAnthropicBedrock,
     AsyncAnthropicVertex,
 )
+from .anthropic_converter import (
+    format_anthropic_response,
+    format_anthropic_input,
+    extract_anthropic_tools,
+    format_anthropic_streaming_content,
+)
 
 __all__ = [
     "Anthropic",
@@ -14,4 +20,8 @@ __all__ = [
     "AsyncAnthropicBedrock",
     "AnthropicVertex",
     "AsyncAnthropicVertex",
+    "format_anthropic_response",
+    "format_anthropic_input",
+    "extract_anthropic_tools",
+    "format_anthropic_streaming_content",
 ]

posthog/ai/anthropic/anthropic.py

@@ -8,14 +8,21 @@ except ImportError:
 
 import time
 import uuid
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional
 
+from posthog.ai.types import StreamingContentBlock, TokenUsage, ToolInProgress
 from posthog.ai.utils import (
     call_llm_and_track_usage,
-    get_model_params,
-    merge_system_prompt,
-    with_privacy_mode,
+    merge_usage_stats,
 )
+from posthog.ai.anthropic.anthropic_converter import (
+    extract_anthropic_usage_from_event,
+    handle_anthropic_content_block_start,
+    handle_anthropic_text_delta,
+    handle_anthropic_tool_delta,
+    finalize_anthropic_tool_input,
+)
+from posthog.ai.sanitization import sanitize_anthropic
 from posthog.client import Client as PostHogClient
 from posthog import setup
 
@@ -61,6 +68,7 @@ class WrappedMessages(Messages):
             posthog_groups: Optional group analytics properties
             **kwargs: Arguments passed to Anthropic's messages.create
         """
+
         if posthog_trace_id is None:
             posthog_trace_id = str(uuid.uuid4())
 
@@ -118,35 +126,66 @@ class WrappedMessages(Messages):
         **kwargs: Any,
     ):
         start_time = time.time()
-        usage_stats: Dict[str, int] = {"input_tokens": 0, "output_tokens": 0}
-        accumulated_content = []
+        usage_stats: TokenUsage = TokenUsage(input_tokens=0, output_tokens=0)
+        accumulated_content = ""
+        content_blocks: List[StreamingContentBlock] = []
+        tools_in_progress: Dict[str, ToolInProgress] = {}
+        current_text_block: Optional[StreamingContentBlock] = None
         response = super().create(**kwargs)
 
         def generator():
             nonlocal usage_stats
-            nonlocal accumulated_content  # noqa: F824
+            nonlocal accumulated_content
+            nonlocal content_blocks
+            nonlocal tools_in_progress
+            nonlocal current_text_block
+
             try:
                 for event in response:
-                    if hasattr(event, "usage") and event.usage:
-                        usage_stats = {
-                            k: getattr(event.usage, k, 0)
-                            for k in [
-                                "input_tokens",
-                                "output_tokens",
-                                "cache_read_input_tokens",
-                                "cache_creation_input_tokens",
-                            ]
-                        }
+                    # Extract usage stats from event
+                    event_usage = extract_anthropic_usage_from_event(event)
+                    merge_usage_stats(usage_stats, event_usage)
 
-                    if hasattr(event, "content") and event.content:
-                        accumulated_content.append(event.content)
+                    # Handle content block start events
+                    if hasattr(event, "type") and event.type == "content_block_start":
+                        block, tool = handle_anthropic_content_block_start(event)
+
+                        if block:
+                            content_blocks.append(block)
+
+                            if block.get("type") == "text":
+                                current_text_block = block
+                            else:
+                                current_text_block = None
+
+                        if tool:
+                            tool_id = tool["block"].get("id")
+                            if tool_id:
+                                tools_in_progress[tool_id] = tool
+
+                    # Handle text delta events
+                    delta_text = handle_anthropic_text_delta(event, current_text_block)
+
+                    if delta_text:
+                        accumulated_content += delta_text
+
+                    # Handle tool input delta events
+                    handle_anthropic_tool_delta(
+                        event, content_blocks, tools_in_progress
+                    )
+
+                    # Handle content block stop events
+                    if hasattr(event, "type") and event.type == "content_block_stop":
+                        current_text_block = None
+                        finalize_anthropic_tool_input(
+                            event, content_blocks, tools_in_progress
+                        )
 
                     yield event
 
             finally:
                 end_time = time.time()
                 latency = end_time - start_time
-                output = "".join(accumulated_content)
 
                 self._capture_streaming_event(
                     posthog_distinct_id,
@@ -157,7 +196,8 @@ class WrappedMessages(Messages):
                     kwargs,
                     usage_stats,
                     latency,
-                    output,
+                    content_blocks,
+                    accumulated_content,
                 )
 
         return generator()
@@ -170,49 +210,39 @@ class WrappedMessages(Messages):
         posthog_privacy_mode: bool,
         posthog_groups: Optional[Dict[str, Any]],
         kwargs: Dict[str, Any],
-        usage_stats: Dict[str, int],
+        usage_stats: TokenUsage,
         latency: float,
-        output: str,
+        content_blocks: List[StreamingContentBlock],
+        accumulated_content: str,
     ):
-        if posthog_trace_id is None:
-            posthog_trace_id = str(uuid.uuid4())
+        from posthog.ai.types import StreamingEventData
+        from posthog.ai.anthropic.anthropic_converter import (
+            format_anthropic_streaming_input,
+            format_anthropic_streaming_output_complete,
+        )
+        from posthog.ai.utils import capture_streaming_event
 
-        event_properties = {
-            "$ai_provider": "anthropic",
-            "$ai_model": kwargs.get("model"),
-            "$ai_model_parameters": get_model_params(kwargs),
-            "$ai_input": with_privacy_mode(
-                self._client._ph_client,
-                posthog_privacy_mode,
-                merge_system_prompt(kwargs, "anthropic"),
-            ),
-            "$ai_output_choices": with_privacy_mode(
-                self._client._ph_client,
-                posthog_privacy_mode,
-                [{"content": output, "role": "assistant"}],
-            ),
-            "$ai_http_status": 200,
-            "$ai_input_tokens": usage_stats.get("input_tokens", 0),
-            "$ai_output_tokens": usage_stats.get("output_tokens", 0),
-            "$ai_cache_read_input_tokens": usage_stats.get(
-                "cache_read_input_tokens", 0
-            ),
-            "$ai_cache_creation_input_tokens": usage_stats.get(
-                "cache_creation_input_tokens", 0
-            ),
-            "$ai_latency": latency,
-            "$ai_trace_id": posthog_trace_id,
-            "$ai_base_url": str(self._client.base_url),
-            **(posthog_properties or {}),
-        }
+        # Prepare standardized event data
+        formatted_input = format_anthropic_streaming_input(kwargs)
+        sanitized_input = sanitize_anthropic(formatted_input)
 
-        if posthog_distinct_id is None:
-            event_properties["$process_person_profile"] = False
+        event_data = StreamingEventData(
+            provider="anthropic",
+            model=kwargs.get("model", "unknown"),
+            base_url=str(self._client.base_url),
+            kwargs=kwargs,
+            formatted_input=sanitized_input,
+            formatted_output=format_anthropic_streaming_output_complete(
+                content_blocks, accumulated_content
+            ),
+            usage_stats=usage_stats,
+            latency=latency,
+            distinct_id=posthog_distinct_id,
+            trace_id=posthog_trace_id,
+            properties=posthog_properties,
+            privacy_mode=posthog_privacy_mode,
+            groups=posthog_groups,
+        )
 
-        if hasattr(self._client._ph_client, "capture"):
-            self._client._ph_client.capture(
-                distinct_id=posthog_distinct_id or posthog_trace_id,
-                event="$ai_generation",
-                properties=event_properties,
-                groups=posthog_groups,
-            )
+        # Use the common capture function
+        capture_streaming_event(self._client._ph_client, event_data)

posthog/ai/anthropic/anthropic_async.py

@@ -8,15 +8,27 @@ except ImportError:
 
 import time
 import uuid
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional
 
 from posthog import setup
+from posthog.ai.types import StreamingContentBlock, TokenUsage, ToolInProgress
 from posthog.ai.utils import (
     call_llm_and_track_usage_async,
+    extract_available_tool_calls,
     get_model_params,
     merge_system_prompt,
+    merge_usage_stats,
     with_privacy_mode,
 )
+from posthog.ai.anthropic.anthropic_converter import (
+    format_anthropic_streaming_content,
+    extract_anthropic_usage_from_event,
+    handle_anthropic_content_block_start,
+    handle_anthropic_text_delta,
+    handle_anthropic_tool_delta,
+    finalize_anthropic_tool_input,
+)
+from posthog.ai.sanitization import sanitize_anthropic
 from posthog.client import Client as PostHogClient
 
 
@@ -61,6 +73,7 @@ class AsyncWrappedMessages(AsyncMessages):
             posthog_groups: Optional group analytics properties
             **kwargs: Arguments passed to Anthropic's messages.create
         """
+
         if posthog_trace_id is None:
             posthog_trace_id = str(uuid.uuid4())
 
@@ -118,35 +131,66 @@ class AsyncWrappedMessages(AsyncMessages):
         **kwargs: Any,
     ):
         start_time = time.time()
-        usage_stats: Dict[str, int] = {"input_tokens": 0, "output_tokens": 0}
-        accumulated_content = []
+        usage_stats: TokenUsage = TokenUsage(input_tokens=0, output_tokens=0)
+        accumulated_content = ""
+        content_blocks: List[StreamingContentBlock] = []
+        tools_in_progress: Dict[str, ToolInProgress] = {}
+        current_text_block: Optional[StreamingContentBlock] = None
         response = await super().create(**kwargs)
 
         async def generator():
             nonlocal usage_stats
-            nonlocal accumulated_content  # noqa: F824
+            nonlocal accumulated_content
+            nonlocal content_blocks
+            nonlocal tools_in_progress
+            nonlocal current_text_block
+
             try:
                 async for event in response:
-                    if hasattr(event, "usage") and event.usage:
-                        usage_stats = {
-                            k: getattr(event.usage, k, 0)
-                            for k in [
-                                "input_tokens",
-                                "output_tokens",
-                                "cache_read_input_tokens",
-                                "cache_creation_input_tokens",
-                            ]
-                        }
+                    # Extract usage stats from event
+                    event_usage = extract_anthropic_usage_from_event(event)
+                    merge_usage_stats(usage_stats, event_usage)
 
-                    if hasattr(event, "content") and event.content:
-                        accumulated_content.append(event.content)
+                    # Handle content block start events
+                    if hasattr(event, "type") and event.type == "content_block_start":
+                        block, tool = handle_anthropic_content_block_start(event)
+
+                        if block:
+                            content_blocks.append(block)
+
+                            if block.get("type") == "text":
+                                current_text_block = block
+                            else:
+                                current_text_block = None
+
+                        if tool:
+                            tool_id = tool["block"].get("id")
+                            if tool_id:
+                                tools_in_progress[tool_id] = tool
+
+                    # Handle text delta events
+                    delta_text = handle_anthropic_text_delta(event, current_text_block)
+
+                    if delta_text:
+                        accumulated_content += delta_text
+
+                    # Handle tool input delta events
+                    handle_anthropic_tool_delta(
+                        event, content_blocks, tools_in_progress
+                    )
+
+                    # Handle content block stop events
+                    if hasattr(event, "type") and event.type == "content_block_stop":
+                        current_text_block = None
+                        finalize_anthropic_tool_input(
+                            event, content_blocks, tools_in_progress
+                        )
 
                     yield event
 
             finally:
                 end_time = time.time()
                 latency = end_time - start_time
-                output = "".join(accumulated_content)
 
                 await self._capture_streaming_event(
                     posthog_distinct_id,
@@ -157,7 +201,8 @@ class AsyncWrappedMessages(AsyncMessages):
                     kwargs,
                     usage_stats,
                     latency,
-                    output,
+                    content_blocks,
+                    accumulated_content,
                 )
 
         return generator()
@@ -170,13 +215,29 @@ class AsyncWrappedMessages(AsyncMessages):
         posthog_privacy_mode: bool,
         posthog_groups: Optional[Dict[str, Any]],
         kwargs: Dict[str, Any],
-        usage_stats: Dict[str, int],
+        usage_stats: TokenUsage,
         latency: float,
-        output: str,
+        content_blocks: List[StreamingContentBlock],
+        accumulated_content: str,
     ):
         if posthog_trace_id is None:
             posthog_trace_id = str(uuid.uuid4())
 
+        # Format output using converter
+        formatted_content = format_anthropic_streaming_content(content_blocks)
+        formatted_output = []
+
+        if formatted_content:
+            formatted_output = [{"role": "assistant", "content": formatted_content}]
+        else:
+            # Fallback to accumulated content if no blocks
+            formatted_output = [
+                {
+                    "role": "assistant",
+                    "content": [{"type": "text", "text": accumulated_content}],
+                }
+            ]
+
         event_properties = {
             "$ai_provider": "anthropic",
             "$ai_model": kwargs.get("model"),
@@ -184,12 +245,12 @@ class AsyncWrappedMessages(AsyncMessages):
             "$ai_input": with_privacy_mode(
                 self._client._ph_client,
                 posthog_privacy_mode,
-                merge_system_prompt(kwargs, "anthropic"),
+                sanitize_anthropic(merge_system_prompt(kwargs, "anthropic")),
             ),
             "$ai_output_choices": with_privacy_mode(
                 self._client._ph_client,
                 posthog_privacy_mode,
-                [{"content": output, "role": "assistant"}],
+                formatted_output,
             ),
             "$ai_http_status": 200,
             "$ai_input_tokens": usage_stats.get("input_tokens", 0),
@@ -206,6 +267,12 @@ class AsyncWrappedMessages(AsyncMessages):
             **(posthog_properties or {}),
         }
 
+        # Add tools if available
+        available_tools = extract_available_tool_calls("anthropic", kwargs)
+
+        if available_tools:
+            event_properties["$ai_tools"] = available_tools
+
         if posthog_distinct_id is None:
             event_properties["$process_person_profile"] = False
 

posthog/ai/anthropic/anthropic_converter.py

@@ -0,0 +1,403 @@
+"""
+Anthropic-specific conversion utilities.
+
+This module handles the conversion of Anthropic API responses and inputs
+into standardized formats for PostHog tracking.
+"""
+
+import json
+from typing import Any, Dict, List, Optional, Tuple
+
+from posthog.ai.types import (
+    FormattedContentItem,
+    FormattedFunctionCall,
+    FormattedMessage,
+    FormattedTextContent,
+    StreamingContentBlock,
+    TokenUsage,
+    ToolInProgress,
+)
+
+
+def format_anthropic_response(response: Any) -> List[FormattedMessage]:
+    """
+    Format an Anthropic response into standardized message format.
+
+    Args:
+        response: The response object from Anthropic API
+
+    Returns:
+        List of formatted messages with role and content
+    """
+
+    output: List[FormattedMessage] = []
+
+    if response is None:
+        return output
+
+    content: List[FormattedContentItem] = []
+
+    # Process content blocks from the response
+    if hasattr(response, "content"):
+        for choice in response.content:
+            if (
+                hasattr(choice, "type")
+                and choice.type == "text"
+                and hasattr(choice, "text")
+                and choice.text
+            ):
+                text_content: FormattedTextContent = {
+                    "type": "text",
+                    "text": choice.text,
+                }
+                content.append(text_content)
+
+            elif (
+                hasattr(choice, "type")
+                and choice.type == "tool_use"
+                and hasattr(choice, "name")
+                and hasattr(choice, "id")
+            ):
+                function_call: FormattedFunctionCall = {
+                    "type": "function",
+                    "id": choice.id,
+                    "function": {
+                        "name": choice.name,
+                        "arguments": getattr(choice, "input", {}),
+                    },
+                }
+                content.append(function_call)
+
+    if content:
+        message: FormattedMessage = {
+            "role": "assistant",
+            "content": content,
+        }
+        output.append(message)
+
+    return output
+
+
+def format_anthropic_input(
+    messages: List[Dict[str, Any]], system: Optional[str] = None
+) -> List[FormattedMessage]:
+    """
+    Format Anthropic input messages with optional system prompt.
+
+    Args:
+        messages: List of message dictionaries
+        system: Optional system prompt to prepend
+
+    Returns:
+        List of formatted messages
+    """
+
+    formatted_messages: List[FormattedMessage] = []
+
+    # Add system message if provided
+    if system is not None:
+        formatted_messages.append({"role": "system", "content": system})
+
+    # Add user messages
+    if messages:
+        for msg in messages:
+            # Messages are already in the correct format, just ensure type safety
+            formatted_msg: FormattedMessage = {
+                "role": msg.get("role", "user"),
+                "content": msg.get("content", ""),
+            }
+            formatted_messages.append(formatted_msg)
+
+    return formatted_messages
+
+
+def extract_anthropic_tools(kwargs: Dict[str, Any]) -> Optional[Any]:
+    """
+    Extract tool definitions from Anthropic API kwargs.
+
+    Args:
+        kwargs: Keyword arguments passed to Anthropic API
+
+    Returns:
+        Tool definitions if present, None otherwise
+    """
+
+    return kwargs.get("tools", None)
+
+
+def format_anthropic_streaming_content(
+    content_blocks: List[StreamingContentBlock],
+) -> List[FormattedContentItem]:
+    """
+    Format content blocks from Anthropic streaming response.
+
+    Used by streaming handlers to format accumulated content blocks.
+
+    Args:
+        content_blocks: List of content block dictionaries from streaming
+
+    Returns:
+        List of formatted content items
+    """
+
+    formatted: List[FormattedContentItem] = []
+
+    for block in content_blocks:
+        if block.get("type") == "text":
+            formatted.append(
+                {
+                    "type": "text",
+                    "text": block.get("text") or "",
+                }
+            )
+
+        elif block.get("type") == "function":
+            formatted.append(
+                {
+                    "type": "function",
+                    "id": block.get("id"),
+                    "function": block.get("function") or {},
+                }
+            )
+
+    return formatted
+
+
+def extract_anthropic_usage_from_response(response: Any) -> TokenUsage:
+    """
+    Extract usage from a full Anthropic response (non-streaming).
+
+    Args:
+        response: The complete response from Anthropic API
+
+    Returns:
+        TokenUsage with standardized usage
+    """
+    if not hasattr(response, "usage"):
+        return TokenUsage(input_tokens=0, output_tokens=0)
+
+    result = TokenUsage(
+        input_tokens=getattr(response.usage, "input_tokens", 0),
+        output_tokens=getattr(response.usage, "output_tokens", 0),
+    )
+
+    if hasattr(response.usage, "cache_read_input_tokens"):
+        cache_read = response.usage.cache_read_input_tokens
+        if cache_read and cache_read > 0:
+            result["cache_read_input_tokens"] = cache_read
+
+    if hasattr(response.usage, "cache_creation_input_tokens"):
+        cache_creation = response.usage.cache_creation_input_tokens
+        if cache_creation and cache_creation > 0:
+            result["cache_creation_input_tokens"] = cache_creation
+
+    return result
+
+
+def extract_anthropic_usage_from_event(event: Any) -> TokenUsage:
+    """
+    Extract usage statistics from an Anthropic streaming event.
+
+    Args:
+        event: Streaming event from Anthropic API
+
+    Returns:
+        Dictionary of usage statistics
+    """
+
+    usage: TokenUsage = TokenUsage()
+
+    # Handle usage stats from message_start event
+    if hasattr(event, "type") and event.type == "message_start":
+        if hasattr(event, "message") and hasattr(event.message, "usage"):
+            usage["input_tokens"] = getattr(event.message.usage, "input_tokens", 0)
+            usage["cache_creation_input_tokens"] = getattr(
+                event.message.usage, "cache_creation_input_tokens", 0
+            )
+            usage["cache_read_input_tokens"] = getattr(
+                event.message.usage, "cache_read_input_tokens", 0
+            )
+
+    # Handle usage stats from message_delta event
+    if hasattr(event, "usage") and event.usage:
+        usage["output_tokens"] = getattr(event.usage, "output_tokens", 0)
+
+    return usage
+
+
+def handle_anthropic_content_block_start(
+    event: Any,
+) -> Tuple[Optional[StreamingContentBlock], Optional[ToolInProgress]]:
+    """
+    Handle content block start event from Anthropic streaming.
+
+    Args:
+        event: Content block start event
+
+    Returns:
+        Tuple of (content_block, tool_in_progress)
+    """
+
+    if not (hasattr(event, "type") and event.type == "content_block_start"):
+        return None, None
+
+    if not hasattr(event, "content_block"):
+        return None, None
+
+    block = event.content_block
+
+    if not hasattr(block, "type"):
+        return None, None
+
+    if block.type == "text":
+        content_block: StreamingContentBlock = {"type": "text", "text": ""}
+        return content_block, None
+
+    elif block.type == "tool_use":
+        tool_block: StreamingContentBlock = {
+            "type": "function",
+            "id": getattr(block, "id", ""),
+            "function": {"name": getattr(block, "name", ""), "arguments": {}},
+        }
+        tool_in_progress: ToolInProgress = {"block": tool_block, "input_string": ""}
+        return tool_block, tool_in_progress
+
+    return None, None
+
+
+def handle_anthropic_text_delta(
+    event: Any, current_block: Optional[StreamingContentBlock]
+) -> Optional[str]:
+    """
+    Handle text delta event from Anthropic streaming.
+
+    Args:
+        event: Delta event
+        current_block: Current text block being accumulated
+
+    Returns:
+        Text delta if present
+    """
+
+    if hasattr(event, "delta") and hasattr(event.delta, "text"):
+        delta_text = event.delta.text or ""
+
+        if current_block is not None and current_block.get("type") == "text":
+            text_val = current_block.get("text")
+            if text_val is not None:
+                current_block["text"] = text_val + delta_text
+            else:
+                current_block["text"] = delta_text
+
+        return delta_text
+
+    return None
+
+
+def handle_anthropic_tool_delta(
+    event: Any,
+    content_blocks: List[StreamingContentBlock],
+    tools_in_progress: Dict[str, ToolInProgress],
+) -> None:
+    """
+    Handle tool input delta event from Anthropic streaming.
+
+    Args:
+        event: Tool delta event
+        content_blocks: List of content blocks
+        tools_in_progress: Dictionary tracking tools being accumulated
+    """
+
+    if not (hasattr(event, "type") and event.type == "content_block_delta"):
+        return
+
+    if not (
+        hasattr(event, "delta")
+        and hasattr(event.delta, "type")
+        and event.delta.type == "input_json_delta"
+    ):
+        return
+
+    if hasattr(event, "index") and event.index < len(content_blocks):
+        block = content_blocks[event.index]
+
+        if block.get("type") == "function" and block.get("id") in tools_in_progress:
+            tool = tools_in_progress[block["id"]]
+            partial_json = getattr(event.delta, "partial_json", "")
+            tool["input_string"] += partial_json
+
+
+def finalize_anthropic_tool_input(
+    event: Any,
+    content_blocks: List[StreamingContentBlock],
+    tools_in_progress: Dict[str, ToolInProgress],
+) -> None:
+    """
+    Finalize tool input when content block stops.
+
+    Args:
+        event: Content block stop event
+        content_blocks: List of content blocks
+        tools_in_progress: Dictionary tracking tools being accumulated
+    """
+
+    if not (hasattr(event, "type") and event.type == "content_block_stop"):
+        return
+
+    if hasattr(event, "index") and event.index < len(content_blocks):
+        block = content_blocks[event.index]
+
+        if block.get("type") == "function" and block.get("id") in tools_in_progress:
+            tool = tools_in_progress[block["id"]]
+
+            try:
+                block["function"]["arguments"] = json.loads(tool["input_string"])
+            except (json.JSONDecodeError, Exception):
+                # Keep empty dict if parsing fails
+                pass
+
+            del tools_in_progress[block["id"]]
+
+
+def format_anthropic_streaming_input(kwargs: Dict[str, Any]) -> Any:
+    """
+    Format Anthropic streaming input using system prompt merging.
+
+    Args:
+        kwargs: Keyword arguments passed to Anthropic API
+
+    Returns:
+        Formatted input ready for PostHog tracking
+    """
+    from posthog.ai.utils import merge_system_prompt
+
+    return merge_system_prompt(kwargs, "anthropic")
+
+
+def format_anthropic_streaming_output_complete(
+    content_blocks: List[StreamingContentBlock], accumulated_content: str
+) -> List[FormattedMessage]:
+    """
+    Format complete Anthropic streaming output.
+
+    Combines existing logic for formatting content blocks with fallback to accumulated content.
+
+    Args:
+        content_blocks: List of content blocks accumulated during streaming
+        accumulated_content: Raw accumulated text content as fallback
+
+    Returns:
+        Formatted messages ready for PostHog tracking
+    """
+    formatted_content = format_anthropic_streaming_content(content_blocks)
+
+    if formatted_content:
+        return [{"role": "assistant", "content": formatted_content}]
+    else:
+        # Fallback to accumulated content if no blocks
+        return [
+            {
+                "role": "assistant",
+                "content": [{"type": "text", "text": accumulated_content}],
+            }
+        ]

posthog/ai/gemini/__init__.py

@@ -1,4 +1,9 @@
 from .gemini import Client
+from .gemini_converter import (
+    format_gemini_input,
+    format_gemini_response,
+    extract_gemini_tools,
+)
 
 
 # Create a genai-like module for perfect drop-in replacement
@@ -8,4 +13,10 @@ class _GenAI:
 
 genai = _GenAI()
 
-__all__ = ["Client", "genai"]
+__all__ = [
+    "Client",
+    "genai",
+    "format_gemini_input",
+    "format_gemini_response",
+    "extract_gemini_tools",
+]

posthog/ai/gemini/gemini.py

@@ -3,6 +3,9 @@ import time
 import uuid
 from typing import Any, Dict, Optional
 
+from posthog.ai.types import TokenUsage, StreamingEventData
+from posthog.ai.utils import merge_system_prompt
+
 try:
     from google import genai
 except ImportError:
@@ -13,9 +16,15 @@ except ImportError:
 from posthog import setup
 from posthog.ai.utils import (
     call_llm_and_track_usage,
-    get_model_params,
-    with_privacy_mode,
+    capture_streaming_event,
+    merge_usage_stats,
 )
+from posthog.ai.gemini.gemini_converter import (
+    extract_gemini_usage_from_chunk,
+    extract_gemini_content_from_chunk,
+    format_gemini_streaming_output,
+)
+from posthog.ai.sanitization import sanitize_gemini
 from posthog.client import Client as PostHogClient
 
 
@@ -71,6 +80,7 @@ class Client:
             posthog_groups: Default groups for all calls (can be overridden per call)
             **kwargs: Additional arguments (for future compatibility)
         """
+
         self._ph_client = posthog_client or setup()
 
         if self._ph_client is None:
@@ -132,6 +142,7 @@ class Models:
             posthog_groups: Default groups for all calls
             **kwargs: Additional arguments (for future compatibility)
         """
+
         self._ph_client = posthog_client or setup()
 
         if self._ph_client is None:
@@ -149,14 +160,19 @@ class Models:
         # Add Vertex AI parameters if provided
         if vertexai is not None:
             client_args["vertexai"] = vertexai
+
         if credentials is not None:
             client_args["credentials"] = credentials
+
         if project is not None:
             client_args["project"] = project
+
         if location is not None:
             client_args["location"] = location
+
         if debug_config is not None:
             client_args["debug_config"] = debug_config
+
         if http_options is not None:
             client_args["http_options"] = http_options
 
@@ -174,6 +190,7 @@ class Models:
                 raise ValueError(
                     "API key must be provided either as parameter or via GOOGLE_API_KEY/API_KEY environment variable"
                 )
+
             client_args["api_key"] = api_key
 
         self._client = genai.Client(**client_args)
@@ -188,6 +205,7 @@ class Models:
         call_groups: Optional[Dict[str, Any]],
     ):
         """Merge call-level PostHog parameters with client defaults."""
+
         # Use call-level values if provided, otherwise fall back to defaults
         distinct_id = (
             call_distinct_id
@@ -203,6 +221,7 @@ class Models:
 
         # Merge properties: default properties + call properties (call properties override)
         properties = dict(self._default_properties)
+
         if call_properties:
             properties.update(call_properties)
 
@@ -238,6 +257,7 @@ class Models:
             posthog_groups: Group analytics properties (overrides client default)
             **kwargs: Arguments passed to Gemini's generate_content
         """
+
         # Merge PostHog parameters
         distinct_id, trace_id, properties, privacy_mode, groups = (
             self._merge_posthog_params(
@@ -276,7 +296,7 @@ class Models:
         **kwargs: Any,
     ):
         start_time = time.time()
-        usage_stats: Dict[str, int] = {"input_tokens": 0, "output_tokens": 0}
+        usage_stats: TokenUsage = TokenUsage(input_tokens=0, output_tokens=0)
         accumulated_content = []
 
         kwargs_without_stream = {"model": model, "contents": contents, **kwargs}
@@ -287,25 +307,24 @@ class Models:
             nonlocal accumulated_content  # noqa: F824
             try:
                 for chunk in response:
-                    if hasattr(chunk, "usage_metadata") and chunk.usage_metadata:
-                        usage_stats = {
-                            "input_tokens": getattr(
-                                chunk.usage_metadata, "prompt_token_count", 0
-                            ),
-                            "output_tokens": getattr(
-                                chunk.usage_metadata, "candidates_token_count", 0
-                            ),
-                        }
+                    # Extract usage stats from chunk
+                    chunk_usage = extract_gemini_usage_from_chunk(chunk)
 
-                    if hasattr(chunk, "text") and chunk.text:
-                        accumulated_content.append(chunk.text)
+                    if chunk_usage:
+                        # Gemini reports cumulative totals, not incremental values
+                        merge_usage_stats(usage_stats, chunk_usage, mode="cumulative")
+
+                    # Extract content from chunk (now returns content blocks)
+                    content_block = extract_gemini_content_from_chunk(chunk)
+
+                    if content_block is not None:
+                        accumulated_content.append(content_block)
 
                     yield chunk
 
             finally:
                 end_time = time.time()
                 latency = end_time - start_time
-                output = "".join(accumulated_content)
 
                 self._capture_streaming_event(
                     model,
@@ -318,7 +337,7 @@ class Models:
                     kwargs,
                     usage_stats,
                     latency,
-                    output,
+                    accumulated_content,
                 )
 
         return generator()
@@ -333,63 +352,39 @@ class Models:
         privacy_mode: bool,
         groups: Optional[Dict[str, Any]],
         kwargs: Dict[str, Any],
-        usage_stats: Dict[str, int],
+        usage_stats: TokenUsage,
         latency: float,
-        output: str,
+        output: Any,
     ):
-        if trace_id is None:
-            trace_id = str(uuid.uuid4())
+        # Prepare standardized event data
+        formatted_input = self._format_input(contents, **kwargs)
+        sanitized_input = sanitize_gemini(formatted_input)
 
-        event_properties = {
-            "$ai_provider": "gemini",
-            "$ai_model": model,
-            "$ai_model_parameters": get_model_params(kwargs),
-            "$ai_input": with_privacy_mode(
-                self._ph_client,
-                privacy_mode,
-                self._format_input(contents),
-            ),
-            "$ai_output_choices": with_privacy_mode(
-                self._ph_client,
-                privacy_mode,
-                [{"content": output, "role": "assistant"}],
-            ),
-            "$ai_http_status": 200,
-            "$ai_input_tokens": usage_stats.get("input_tokens", 0),
-            "$ai_output_tokens": usage_stats.get("output_tokens", 0),
-            "$ai_latency": latency,
-            "$ai_trace_id": trace_id,
-            "$ai_base_url": self._base_url,
-            **(properties or {}),
-        }
+        event_data = StreamingEventData(
+            provider="gemini",
+            model=model,
+            base_url=self._base_url,
+            kwargs=kwargs,
+            formatted_input=sanitized_input,
+            formatted_output=format_gemini_streaming_output(output),
+            usage_stats=usage_stats,
+            latency=latency,
+            distinct_id=distinct_id,
+            trace_id=trace_id,
+            properties=properties,
+            privacy_mode=privacy_mode,
+            groups=groups,
+        )
 
-        if distinct_id is None:
-            event_properties["$process_person_profile"] = False
+        # Use the common capture function
+        capture_streaming_event(self._ph_client, event_data)
 
-        if hasattr(self._ph_client, "capture"):
-            self._ph_client.capture(
-                distinct_id=distinct_id,
-                event="$ai_generation",
-                properties=event_properties,
-                groups=groups,
-            )
-
-    def _format_input(self, contents):
+    def _format_input(self, contents, **kwargs):
         """Format input contents for PostHog tracking"""
-        if isinstance(contents, str):
-            return [{"role": "user", "content": contents}]
-        elif isinstance(contents, list):
-            formatted = []
-            for item in contents:
-                if isinstance(item, str):
-                    formatted.append({"role": "user", "content": item})
-                elif hasattr(item, "text"):
-                    formatted.append({"role": "user", "content": item.text})
-                else:
-                    formatted.append({"role": "user", "content": str(item)})
-            return formatted
-        else:
-            return [{"role": "user", "content": str(contents)}]
+
+        # Create kwargs dict with contents for merge_system_prompt
+        input_kwargs = {"contents": contents, **kwargs}
+        return merge_system_prompt(input_kwargs, "gemini")
 
     def generate_content_stream(
         self,

posthog/ai/gemini/gemini_converter.py

@@ -0,0 +1,516 @@
+"""
+Gemini-specific conversion utilities.
+
+This module handles the conversion of Gemini API responses and inputs
+into standardized formats for PostHog tracking.
+"""
+
+from typing import Any, Dict, List, Optional, TypedDict, Union
+
+from posthog.ai.types import (
+    FormattedContentItem,
+    FormattedMessage,
+    TokenUsage,
+)
+
+
+class GeminiPart(TypedDict, total=False):
+    """Represents a part in a Gemini message."""
+
+    text: str
+
+
+class GeminiMessage(TypedDict, total=False):
+    """Represents a Gemini message with various possible fields."""
+
+    role: str
+    parts: List[Union[GeminiPart, Dict[str, Any]]]
+    content: Union[str, List[Any]]
+    text: str
+
+
+def _extract_text_from_parts(parts: List[Any]) -> str:
+    """
+    Extract and concatenate text from a parts array.
+
+    Args:
+        parts: List of parts that may contain text content
+
+    Returns:
+        Concatenated text from all parts
+    """
+
+    content_parts = []
+
+    for part in parts:
+        if isinstance(part, dict) and "text" in part:
+            content_parts.append(part["text"])
+
+        elif isinstance(part, str):
+            content_parts.append(part)
+
+        elif hasattr(part, "text"):
+            # Get the text attribute value
+            text_value = getattr(part, "text", "")
+            content_parts.append(text_value if text_value else str(part))
+
+        else:
+            content_parts.append(str(part))
+
+    return "".join(content_parts)
+
+
+def _format_dict_message(item: Dict[str, Any]) -> FormattedMessage:
+    """
+    Format a dictionary message into standardized format.
+
+    Args:
+        item: Dictionary containing message data
+
+    Returns:
+        Formatted message with role and content
+    """
+
+    # Handle dict format with parts array (Gemini-specific format)
+    if "parts" in item and isinstance(item["parts"], list):
+        content = _extract_text_from_parts(item["parts"])
+        return {"role": item.get("role", "user"), "content": content}
+
+    # Handle dict with content field
+    if "content" in item:
+        content = item["content"]
+
+        if isinstance(content, list):
+            # If content is a list, extract text from it
+            content = _extract_text_from_parts(content)
+
+        elif not isinstance(content, str):
+            content = str(content)
+
+        return {"role": item.get("role", "user"), "content": content}
+
+    # Handle dict with text field
+    if "text" in item:
+        return {"role": item.get("role", "user"), "content": item["text"]}
+
+    # Fallback to string representation
+    return {"role": "user", "content": str(item)}
+
+
+def _format_object_message(item: Any) -> FormattedMessage:
+    """
+    Format an object (with attributes) into standardized format.
+
+    Args:
+        item: Object that may have text or parts attributes
+
+    Returns:
+        Formatted message with role and content
+    """
+
+    # Handle object with parts attribute
+    if hasattr(item, "parts") and hasattr(item.parts, "__iter__"):
+        content = _extract_text_from_parts(item.parts)
+        role = getattr(item, "role", "user") if hasattr(item, "role") else "user"
+
+        # Ensure role is a string
+        if not isinstance(role, str):
+            role = "user"
+
+        return {"role": role, "content": content}
+
+    # Handle object with text attribute
+    if hasattr(item, "text"):
+        role = getattr(item, "role", "user") if hasattr(item, "role") else "user"
+
+        # Ensure role is a string
+        if not isinstance(role, str):
+            role = "user"
+
+        return {"role": role, "content": item.text}
+
+    # Handle object with content attribute
+    if hasattr(item, "content"):
+        role = getattr(item, "role", "user") if hasattr(item, "role") else "user"
+
+        # Ensure role is a string
+        if not isinstance(role, str):
+            role = "user"
+
+        content = item.content
+
+        if isinstance(content, list):
+            content = _extract_text_from_parts(content)
+
+        elif not isinstance(content, str):
+            content = str(content)
+        return {"role": role, "content": content}
+
+    # Fallback to string representation
+    return {"role": "user", "content": str(item)}
+
+
+def format_gemini_response(response: Any) -> List[FormattedMessage]:
+    """
+    Format a Gemini response into standardized message format.
+
+    Args:
+        response: The response object from Gemini API
+
+    Returns:
+        List of formatted messages with role and content
+    """
+
+    output: List[FormattedMessage] = []
+
+    if response is None:
+        return output
+
+    if hasattr(response, "candidates") and response.candidates:
+        for candidate in response.candidates:
+            if hasattr(candidate, "content") and candidate.content:
+                content: List[FormattedContentItem] = []
+
+                if hasattr(candidate.content, "parts") and candidate.content.parts:
+                    for part in candidate.content.parts:
+                        if hasattr(part, "text") and part.text:
+                            content.append(
+                                {
+                                    "type": "text",
+                                    "text": part.text,
+                                }
+                            )
+
+                        elif hasattr(part, "function_call") and part.function_call:
+                            function_call = part.function_call
+                            content.append(
+                                {
+                                    "type": "function",
+                                    "function": {
+                                        "name": function_call.name,
+                                        "arguments": function_call.args,
+                                    },
+                                }
+                            )
+
+                if content:
+                    output.append(
+                        {
+                            "role": "assistant",
+                            "content": content,
+                        }
+                    )
+
+            elif hasattr(candidate, "text") and candidate.text:
+                output.append(
+                    {
+                        "role": "assistant",
+                        "content": [{"type": "text", "text": candidate.text}],
+                    }
+                )
+
+    elif hasattr(response, "text") and response.text:
+        output.append(
+            {
+                "role": "assistant",
+                "content": [{"type": "text", "text": response.text}],
+            }
+        )
+
+    return output
+
+
+def extract_gemini_system_instruction(config: Any) -> Optional[str]:
+    """
+    Extract system instruction from Gemini config parameter.
+
+    Args:
+        config: Config object or dict that may contain system instruction
+
+    Returns:
+        System instruction string if present, None otherwise
+    """
+    if config is None:
+        return None
+
+    # Handle different config formats
+    if hasattr(config, "system_instruction"):
+        return config.system_instruction
+    elif isinstance(config, dict) and "system_instruction" in config:
+        return config["system_instruction"]
+    elif isinstance(config, dict) and "systemInstruction" in config:
+        return config["systemInstruction"]
+
+    return None
+
+
+def extract_gemini_tools(kwargs: Dict[str, Any]) -> Optional[Any]:
+    """
+    Extract tool definitions from Gemini API kwargs.
+
+    Args:
+        kwargs: Keyword arguments passed to Gemini API
+
+    Returns:
+        Tool definitions if present, None otherwise
+    """
+
+    if "config" in kwargs and hasattr(kwargs["config"], "tools"):
+        return kwargs["config"].tools
+
+    return None
+
+
+def format_gemini_input_with_system(
+    contents: Any, config: Any = None
+) -> List[FormattedMessage]:
+    """
+    Format Gemini input contents into standardized message format, including system instruction handling.
+
+    Args:
+        contents: Input contents in various possible formats
+        config: Config object or dict that may contain system instruction
+
+    Returns:
+        List of formatted messages with role and content fields, with system message prepended if needed
+    """
+    formatted_messages = format_gemini_input(contents)
+
+    # Check if system instruction is provided in config parameter
+    system_instruction = extract_gemini_system_instruction(config)
+
+    if system_instruction is not None:
+        has_system = any(msg.get("role") == "system" for msg in formatted_messages)
+        if not has_system:
+            from posthog.ai.types import FormattedMessage
+
+            system_message: FormattedMessage = {
+                "role": "system",
+                "content": system_instruction,
+            }
+            formatted_messages = [system_message] + list(formatted_messages)
+
+    return formatted_messages
+
+
+def format_gemini_input(contents: Any) -> List[FormattedMessage]:
+    """
+    Format Gemini input contents into standardized message format for PostHog tracking.
+
+    This function handles various input formats:
+    - String inputs
+    - List of strings, dicts, or objects
+    - Single dict or object
+    - Gemini-specific format with parts array
+
+    Args:
+        contents: Input contents in various possible formats
+
+    Returns:
+        List of formatted messages with role and content fields
+    """
+
+    # Handle string input
+    if isinstance(contents, str):
+        return [{"role": "user", "content": contents}]
+
+    # Handle list input
+    if isinstance(contents, list):
+        formatted: List[FormattedMessage] = []
+
+        for item in contents:
+            if isinstance(item, str):
+                formatted.append({"role": "user", "content": item})
+
+            elif isinstance(item, dict):
+                formatted.append(_format_dict_message(item))
+
+            else:
+                formatted.append(_format_object_message(item))
+
+        return formatted
+
+    # Handle single dict input
+    if isinstance(contents, dict):
+        return [_format_dict_message(contents)]
+
+    # Handle single object input
+    return [_format_object_message(contents)]
+
+
+def _extract_usage_from_metadata(metadata: Any) -> TokenUsage:
+    """
+    Common logic to extract usage from Gemini metadata.
+    Used by both streaming and non-streaming paths.
+
+    Args:
+        metadata: usage_metadata from Gemini response or chunk
+
+    Returns:
+        TokenUsage with standardized usage
+    """
+    usage = TokenUsage(
+        input_tokens=getattr(metadata, "prompt_token_count", 0),
+        output_tokens=getattr(metadata, "candidates_token_count", 0),
+    )
+
+    # Add cache tokens if present (don't add if 0)
+    if hasattr(metadata, "cached_content_token_count"):
+        cache_tokens = metadata.cached_content_token_count
+        if cache_tokens and cache_tokens > 0:
+            usage["cache_read_input_tokens"] = cache_tokens
+
+    # Add reasoning tokens if present (don't add if 0)
+    if hasattr(metadata, "thoughts_token_count"):
+        reasoning_tokens = metadata.thoughts_token_count
+        if reasoning_tokens and reasoning_tokens > 0:
+            usage["reasoning_tokens"] = reasoning_tokens
+
+    return usage
+
+
+def extract_gemini_usage_from_response(response: Any) -> TokenUsage:
+    """
+    Extract usage statistics from a full Gemini response (non-streaming).
+
+    Args:
+        response: The complete response from Gemini API
+
+    Returns:
+        TokenUsage with standardized usage statistics
+    """
+    if not hasattr(response, "usage_metadata") or not response.usage_metadata:
+        return TokenUsage(input_tokens=0, output_tokens=0)
+
+    return _extract_usage_from_metadata(response.usage_metadata)
+
+
+def extract_gemini_usage_from_chunk(chunk: Any) -> TokenUsage:
+    """
+    Extract usage statistics from a Gemini streaming chunk.
+
+    Args:
+        chunk: Streaming chunk from Gemini API
+
+    Returns:
+        TokenUsage with standardized usage statistics
+    """
+
+    usage: TokenUsage = TokenUsage()
+
+    if not hasattr(chunk, "usage_metadata") or not chunk.usage_metadata:
+        return usage
+
+    # Use the shared helper to extract usage
+    usage = _extract_usage_from_metadata(chunk.usage_metadata)
+
+    return usage
+
+
+def extract_gemini_content_from_chunk(chunk: Any) -> Optional[Dict[str, Any]]:
+    """
+    Extract content (text or function call) from a Gemini streaming chunk.
+
+    Args:
+        chunk: Streaming chunk from Gemini API
+
+    Returns:
+        Content block dictionary if present, None otherwise
+    """
+
+    # Check for text content
+    if hasattr(chunk, "text") and chunk.text:
+        return {"type": "text", "text": chunk.text}
+
+    # Check for function calls in candidates
+    if hasattr(chunk, "candidates") and chunk.candidates:
+        for candidate in chunk.candidates:
+            if hasattr(candidate, "content") and candidate.content:
+                if hasattr(candidate.content, "parts") and candidate.content.parts:
+                    for part in candidate.content.parts:
+                        # Check for function_call part
+                        if hasattr(part, "function_call") and part.function_call:
+                            function_call = part.function_call
+                            return {
+                                "type": "function",
+                                "function": {
+                                    "name": function_call.name,
+                                    "arguments": function_call.args,
+                                },
+                            }
+                        # Also check for text in parts
+                        elif hasattr(part, "text") and part.text:
+                            return {"type": "text", "text": part.text}
+
+    return None
+
+
+def format_gemini_streaming_output(
+    accumulated_content: Union[str, List[Any]],
+) -> List[FormattedMessage]:
+    """
+    Format the final output from Gemini streaming.
+
+    Args:
+        accumulated_content: Accumulated content from streaming (string, list of strings, or list of content blocks)
+
+    Returns:
+        List of formatted messages
+    """
+
+    # Handle legacy string input (backward compatibility)
+    if isinstance(accumulated_content, str):
+        return [
+            {
+                "role": "assistant",
+                "content": [{"type": "text", "text": accumulated_content}],
+            }
+        ]
+
+    # Handle list input
+    if isinstance(accumulated_content, list):
+        content: List[FormattedContentItem] = []
+        text_parts = []
+
+        for item in accumulated_content:
+            if isinstance(item, str):
+                # Legacy support: accumulate strings
+                text_parts.append(item)
+            elif isinstance(item, dict):
+                # New format: content blocks
+                if item.get("type") == "text":
+                    text_parts.append(item.get("text", ""))
+                elif item.get("type") == "function":
+                    # If we have accumulated text, add it first
+                    if text_parts:
+                        content.append(
+                            {
+                                "type": "text",
+                                "text": "".join(text_parts),
+                            }
+                        )
+                        text_parts = []
+
+                    # Add the function call
+                    content.append(
+                        {
+                            "type": "function",
+                            "function": item.get("function", {}),
+                        }
+                    )
+
+        # Add any remaining text
+        if text_parts:
+            content.append(
+                {
+                    "type": "text",
+                    "text": "".join(text_parts),
+                }
+            )
+
+        # If we have content, return it
+        if content:
+            return [{"role": "assistant", "content": content}]
+
+    # Fallback for empty or unexpected input
+    return [{"role": "assistant", "content": [{"type": "text", "text": ""}]}]

posthog/ai/langchain/callbacks.py

@@ -37,6 +37,7 @@ from pydantic import BaseModel
 
 from posthog import setup
 from posthog.ai.utils import get_model_params, with_privacy_mode
+from posthog.ai.sanitization import sanitize_langchain
 from posthog.client import Client
 
 log = logging.getLogger("posthog")
@@ -480,11 +481,12 @@ class CallbackHandler(BaseCallbackHandler):
         event_properties = {
             "$ai_trace_id": trace_id,
             "$ai_input_state": with_privacy_mode(
-                self._ph_client, self._privacy_mode, run.input
+                self._ph_client, self._privacy_mode, sanitize_langchain(run.input)
             ),
             "$ai_latency": run.latency,
             "$ai_span_name": run.name,
             "$ai_span_id": run_id,
+            "$ai_framework": "langchain",
         }
         if parent_run_id is not None:
             event_properties["$ai_parent_id"] = parent_run_id
@@ -550,11 +552,12 @@ class CallbackHandler(BaseCallbackHandler):
             "$ai_model": run.model,
             "$ai_model_parameters": run.model_params,
             "$ai_input": with_privacy_mode(
-                self._ph_client, self._privacy_mode, run.input
+                self._ph_client, self._privacy_mode, sanitize_langchain(run.input)
             ),
             "$ai_http_status": 200,
             "$ai_latency": run.latency,
             "$ai_base_url": run.base_url,
+            "$ai_framework": "langchain",
         }
 
         if run.tools:

posthog/ai/openai/__init__.py

@@ -1,5 +1,20 @@
 from .openai import OpenAI
 from .openai_async import AsyncOpenAI
 from .openai_providers import AsyncAzureOpenAI, AzureOpenAI
+from .openai_converter import (
+    format_openai_response,
+    format_openai_input,
+    extract_openai_tools,
+    format_openai_streaming_content,
+)
 
-__all__ = ["OpenAI", "AsyncOpenAI", "AzureOpenAI", "AsyncAzureOpenAI"]
+__all__ = [
+    "OpenAI",
+    "AsyncOpenAI",
+    "AzureOpenAI",
+    "AsyncAzureOpenAI",
+    "format_openai_response",
+    "format_openai_input",
+    "extract_openai_tools",
+    "format_openai_streaming_content",
+]

posthog/ai/openai/openai.py

@@ -2,6 +2,8 @@ import time
 import uuid
 from typing import Any, Dict, List, Optional
 
+from posthog.ai.types import TokenUsage
+
 try:
     import openai
 except ImportError:
@@ -12,9 +14,16 @@ except ImportError:
 from posthog.ai.utils import (
     call_llm_and_track_usage,
     extract_available_tool_calls,
-    get_model_params,
+    merge_usage_stats,
     with_privacy_mode,
 )
+from posthog.ai.openai.openai_converter import (
+    extract_openai_usage_from_chunk,
+    extract_openai_content_from_chunk,
+    extract_openai_tool_calls_from_chunk,
+    accumulate_openai_tool_calls,
+)
+from posthog.ai.sanitization import sanitize_openai, sanitize_openai_response
 from posthog.client import Client as PostHogClient
 from posthog import setup
 
@@ -33,6 +42,7 @@ class OpenAI(openai.OpenAI):
             posthog_client: If provided, events will be captured via this client instead of the global `posthog`.
             **openai_config: Any additional keyword args to set on openai (e.g. organization="xxx").
         """
+
         super().__init__(**kwargs)
         self._ph_client = posthog_client or setup()
 
@@ -112,7 +122,7 @@ class WrappedResponses:
         **kwargs: Any,
     ):
         start_time = time.time()
-        usage_stats: Dict[str, int] = {}
+        usage_stats: TokenUsage = TokenUsage()
         final_content = []
         response = self._original.create(**kwargs)
 
@@ -122,35 +132,17 @@ class WrappedResponses:
 
             try:
                 for chunk in response:
-                    if hasattr(chunk, "type") and chunk.type == "response.completed":
-                        res = chunk.response
-                        if res.output and len(res.output) > 0:
-                            final_content.append(res.output[0])
+                    # Extract usage stats from chunk
+                    chunk_usage = extract_openai_usage_from_chunk(chunk, "responses")
 
-                    if hasattr(chunk, "usage") and chunk.usage:
-                        usage_stats = {
-                            k: getattr(chunk.usage, k, 0)
-                            for k in [
-                                "input_tokens",
-                                "output_tokens",
-                                "total_tokens",
-                            ]
-                        }
+                    if chunk_usage:
+                        merge_usage_stats(usage_stats, chunk_usage)
 
-                        # Add support for cached tokens
-                        if hasattr(chunk.usage, "output_tokens_details") and hasattr(
-                            chunk.usage.output_tokens_details, "reasoning_tokens"
-                        ):
-                            usage_stats["reasoning_tokens"] = (
-                                chunk.usage.output_tokens_details.reasoning_tokens
-                            )
+                    # Extract content from chunk
+                    content = extract_openai_content_from_chunk(chunk, "responses")
 
-                        if hasattr(chunk.usage, "input_tokens_details") and hasattr(
-                            chunk.usage.input_tokens_details, "cached_tokens"
-                        ):
-                            usage_stats["cache_read_input_tokens"] = (
-                                chunk.usage.input_tokens_details.cached_tokens
-                            )
+                    if content is not None:
+                        final_content.append(content)
 
                     yield chunk
 
@@ -168,7 +160,7 @@ class WrappedResponses:
                     usage_stats,
                     latency,
                     output,
-                    extract_available_tool_calls("openai", kwargs),
+                    None,  # Responses API doesn't have tools
                 )
 
         return generator()
@@ -181,52 +173,40 @@ class WrappedResponses:
         posthog_privacy_mode: bool,
         posthog_groups: Optional[Dict[str, Any]],
         kwargs: Dict[str, Any],
-        usage_stats: Dict[str, int],
+        usage_stats: TokenUsage,
         latency: float,
         output: Any,
         available_tool_calls: Optional[List[Dict[str, Any]]] = None,
     ):
-        if posthog_trace_id is None:
-            posthog_trace_id = str(uuid.uuid4())
+        from posthog.ai.types import StreamingEventData
+        from posthog.ai.openai.openai_converter import (
+            format_openai_streaming_input,
+            format_openai_streaming_output,
+        )
+        from posthog.ai.utils import capture_streaming_event
 
-        event_properties = {
-            "$ai_provider": "openai",
-            "$ai_model": kwargs.get("model"),
-            "$ai_model_parameters": get_model_params(kwargs),
-            "$ai_input": with_privacy_mode(
-                self._client._ph_client, posthog_privacy_mode, kwargs.get("input")
-            ),
-            "$ai_output_choices": with_privacy_mode(
-                self._client._ph_client,
-                posthog_privacy_mode,
-                output,
-            ),
-            "$ai_http_status": 200,
-            "$ai_input_tokens": usage_stats.get("input_tokens", 0),
-            "$ai_output_tokens": usage_stats.get("output_tokens", 0),
-            "$ai_cache_read_input_tokens": usage_stats.get(
-                "cache_read_input_tokens", 0
-            ),
-            "$ai_reasoning_tokens": usage_stats.get("reasoning_tokens", 0),
-            "$ai_latency": latency,
-            "$ai_trace_id": posthog_trace_id,
-            "$ai_base_url": str(self._client.base_url),
-            **(posthog_properties or {}),
-        }
+        # Prepare standardized event data
+        formatted_input = format_openai_streaming_input(kwargs, "responses")
+        sanitized_input = sanitize_openai_response(formatted_input)
 
-        if available_tool_calls:
-            event_properties["$ai_tools"] = available_tool_calls
+        event_data = StreamingEventData(
+            provider="openai",
+            model=kwargs.get("model", "unknown"),
+            base_url=str(self._client.base_url),
+            kwargs=kwargs,
+            formatted_input=sanitized_input,
+            formatted_output=format_openai_streaming_output(output, "responses"),
+            usage_stats=usage_stats,
+            latency=latency,
+            distinct_id=posthog_distinct_id,
+            trace_id=posthog_trace_id,
+            properties=posthog_properties,
+            privacy_mode=posthog_privacy_mode,
+            groups=posthog_groups,
+        )
 
-        if posthog_distinct_id is None:
-            event_properties["$process_person_profile"] = False
-
-        if hasattr(self._client._ph_client, "capture"):
-            self._client._ph_client.capture(
-                distinct_id=posthog_distinct_id or posthog_trace_id,
-                event="$ai_generation",
-                properties=event_properties,
-                groups=posthog_groups,
-            )
+        # Use the common capture function
+        capture_streaming_event(self._client._ph_client, event_data)
 
     def parse(
         self,
@@ -337,8 +317,9 @@ class WrappedCompletions:
         **kwargs: Any,
     ):
         start_time = time.time()
-        usage_stats: Dict[str, int] = {}
+        usage_stats: TokenUsage = TokenUsage()
         accumulated_content = []
+        accumulated_tool_calls: Dict[int, Dict[str, Any]] = {}
         if "stream_options" not in kwargs:
             kwargs["stream_options"] = {}
         kwargs["stream_options"]["include_usage"] = True
@@ -347,50 +328,42 @@ class WrappedCompletions:
         def generator():
             nonlocal usage_stats
             nonlocal accumulated_content  # noqa: F824
+            nonlocal accumulated_tool_calls
 
             try:
                 for chunk in response:
-                    if hasattr(chunk, "usage") and chunk.usage:
-                        usage_stats = {
-                            k: getattr(chunk.usage, k, 0)
-                            for k in [
-                                "prompt_tokens",
-                                "completion_tokens",
-                                "total_tokens",
-                            ]
-                        }
+                    # Extract usage stats from chunk
+                    chunk_usage = extract_openai_usage_from_chunk(chunk, "chat")
 
-                        # Add support for cached tokens
-                        if hasattr(chunk.usage, "prompt_tokens_details") and hasattr(
-                            chunk.usage.prompt_tokens_details, "cached_tokens"
-                        ):
-                            usage_stats["cache_read_input_tokens"] = (
-                                chunk.usage.prompt_tokens_details.cached_tokens
-                            )
+                    if chunk_usage:
+                        merge_usage_stats(usage_stats, chunk_usage)
 
-                        if hasattr(chunk.usage, "output_tokens_details") and hasattr(
-                            chunk.usage.output_tokens_details, "reasoning_tokens"
-                        ):
-                            usage_stats["reasoning_tokens"] = (
-                                chunk.usage.output_tokens_details.reasoning_tokens
-                            )
+                    # Extract content from chunk
+                    content = extract_openai_content_from_chunk(chunk, "chat")
 
-                    if (
-                        hasattr(chunk, "choices")
-                        and chunk.choices
-                        and len(chunk.choices) > 0
-                    ):
-                        if chunk.choices[0].delta and chunk.choices[0].delta.content:
-                            content = chunk.choices[0].delta.content
-                            if content:
-                                accumulated_content.append(content)
+                    if content is not None:
+                        accumulated_content.append(content)
+
+                    # Extract and accumulate tool calls from chunk
+                    chunk_tool_calls = extract_openai_tool_calls_from_chunk(chunk)
+                    if chunk_tool_calls:
+                        accumulate_openai_tool_calls(
+                            accumulated_tool_calls, chunk_tool_calls
+                        )
 
                     yield chunk
 
             finally:
                 end_time = time.time()
                 latency = end_time - start_time
-                output = "".join(accumulated_content)
+
+                # Convert accumulated tool calls dict to list
+                tool_calls_list = (
+                    list(accumulated_tool_calls.values())
+                    if accumulated_tool_calls
+                    else None
+                )
+
                 self._capture_streaming_event(
                     posthog_distinct_id,
                     posthog_trace_id,
@@ -400,7 +373,8 @@ class WrappedCompletions:
                     kwargs,
                     usage_stats,
                     latency,
-                    output,
+                    accumulated_content,
+                    tool_calls_list,
                     extract_available_tool_calls("openai", kwargs),
                 )
 
@@ -414,52 +388,41 @@ class WrappedCompletions:
         posthog_privacy_mode: bool,
         posthog_groups: Optional[Dict[str, Any]],
         kwargs: Dict[str, Any],
-        usage_stats: Dict[str, int],
+        usage_stats: TokenUsage,
         latency: float,
         output: Any,
+        tool_calls: Optional[List[Dict[str, Any]]] = None,
         available_tool_calls: Optional[List[Dict[str, Any]]] = None,
     ):
-        if posthog_trace_id is None:
-            posthog_trace_id = str(uuid.uuid4())
+        from posthog.ai.types import StreamingEventData
+        from posthog.ai.openai.openai_converter import (
+            format_openai_streaming_input,
+            format_openai_streaming_output,
+        )
+        from posthog.ai.utils import capture_streaming_event
 
-        event_properties = {
-            "$ai_provider": "openai",
-            "$ai_model": kwargs.get("model"),
-            "$ai_model_parameters": get_model_params(kwargs),
-            "$ai_input": with_privacy_mode(
-                self._client._ph_client, posthog_privacy_mode, kwargs.get("messages")
-            ),
-            "$ai_output_choices": with_privacy_mode(
-                self._client._ph_client,
-                posthog_privacy_mode,
-                [{"content": output, "role": "assistant"}],
-            ),
-            "$ai_http_status": 200,
-            "$ai_input_tokens": usage_stats.get("prompt_tokens", 0),
-            "$ai_output_tokens": usage_stats.get("completion_tokens", 0),
-            "$ai_cache_read_input_tokens": usage_stats.get(
-                "cache_read_input_tokens", 0
-            ),
-            "$ai_reasoning_tokens": usage_stats.get("reasoning_tokens", 0),
-            "$ai_latency": latency,
-            "$ai_trace_id": posthog_trace_id,
-            "$ai_base_url": str(self._client.base_url),
-            **(posthog_properties or {}),
-        }
+        # Prepare standardized event data
+        formatted_input = format_openai_streaming_input(kwargs, "chat")
+        sanitized_input = sanitize_openai(formatted_input)
 
-        if available_tool_calls:
-            event_properties["$ai_tools"] = available_tool_calls
+        event_data = StreamingEventData(
+            provider="openai",
+            model=kwargs.get("model", "unknown"),
+            base_url=str(self._client.base_url),
+            kwargs=kwargs,
+            formatted_input=sanitized_input,
+            formatted_output=format_openai_streaming_output(output, "chat", tool_calls),
+            usage_stats=usage_stats,
+            latency=latency,
+            distinct_id=posthog_distinct_id,
+            trace_id=posthog_trace_id,
+            properties=posthog_properties,
+            privacy_mode=posthog_privacy_mode,
+            groups=posthog_groups,
+        )
 
-        if posthog_distinct_id is None:
-            event_properties["$process_person_profile"] = False
-
-        if hasattr(self._client._ph_client, "capture"):
-            self._client._ph_client.capture(
-                distinct_id=posthog_distinct_id or posthog_trace_id,
-                event="$ai_generation",
-                properties=event_properties,
-                groups=posthog_groups,
-            )
+        # Use the common capture function
+        capture_streaming_event(self._client._ph_client, event_data)
 
 
 class WrappedEmbeddings:
@@ -496,6 +459,7 @@ class WrappedEmbeddings:
         Returns:
             The response from OpenAI's embeddings.create call.
         """
+
         if posthog_trace_id is None:
             posthog_trace_id = str(uuid.uuid4())
 
@@ -518,7 +482,9 @@ class WrappedEmbeddings:
             "$ai_provider": "openai",
             "$ai_model": kwargs.get("model"),
             "$ai_input": with_privacy_mode(
-                self._client._ph_client, posthog_privacy_mode, kwargs.get("input")
+                self._client._ph_client,
+                posthog_privacy_mode,
+                sanitize_openai_response(kwargs.get("input")),
             ),
             "$ai_http_status": 200,
             "$ai_input_tokens": usage_stats.get("prompt_tokens", 0),

posthog/ai/openai/openai_async.py

@@ -2,6 +2,8 @@ import time
 import uuid
 from typing import Any, Dict, List, Optional
 
+from posthog.ai.types import TokenUsage
+
 try:
     import openai
 except ImportError:
@@ -14,8 +16,17 @@ from posthog.ai.utils import (
     call_llm_and_track_usage_async,
     extract_available_tool_calls,
     get_model_params,
+    merge_usage_stats,
     with_privacy_mode,
 )
+from posthog.ai.openai.openai_converter import (
+    extract_openai_usage_from_chunk,
+    extract_openai_content_from_chunk,
+    extract_openai_tool_calls_from_chunk,
+    accumulate_openai_tool_calls,
+    format_openai_streaming_output,
+)
+from posthog.ai.sanitization import sanitize_openai, sanitize_openai_response
 from posthog.client import Client as PostHogClient
 
 
@@ -34,6 +45,7 @@ class AsyncOpenAI(openai.AsyncOpenAI):
                             of the global posthog.
             **openai_config: Any additional keyword args to set on openai (e.g. organization="xxx").
         """
+
         super().__init__(**kwargs)
         self._ph_client = posthog_client or setup()
 
@@ -66,6 +78,7 @@ class WrappedResponses:
 
     def __getattr__(self, name):
         """Fallback to original responses object for any methods we don't explicitly handle."""
+
         return getattr(self._original, name)
 
     async def create(
@@ -113,7 +126,7 @@ class WrappedResponses:
         **kwargs: Any,
     ):
         start_time = time.time()
-        usage_stats: Dict[str, int] = {}
+        usage_stats: TokenUsage = TokenUsage()
         final_content = []
         response = await self._original.create(**kwargs)
 
@@ -123,35 +136,17 @@ class WrappedResponses:
 
             try:
                 async for chunk in response:
-                    if hasattr(chunk, "type") and chunk.type == "response.completed":
-                        res = chunk.response
-                        if res.output and len(res.output) > 0:
-                            final_content.append(res.output[0])
+                    # Extract usage stats from chunk
+                    chunk_usage = extract_openai_usage_from_chunk(chunk, "responses")
 
-                    if hasattr(chunk, "usage") and chunk.usage:
-                        usage_stats = {
-                            k: getattr(chunk.usage, k, 0)
-                            for k in [
-                                "input_tokens",
-                                "output_tokens",
-                                "total_tokens",
-                            ]
-                        }
+                    if chunk_usage:
+                        merge_usage_stats(usage_stats, chunk_usage)
 
-                        # Add support for cached tokens
-                        if hasattr(chunk.usage, "output_tokens_details") and hasattr(
-                            chunk.usage.output_tokens_details, "reasoning_tokens"
-                        ):
-                            usage_stats["reasoning_tokens"] = (
-                                chunk.usage.output_tokens_details.reasoning_tokens
-                            )
+                    # Extract content from chunk
+                    content = extract_openai_content_from_chunk(chunk, "responses")
 
-                        if hasattr(chunk.usage, "input_tokens_details") and hasattr(
-                            chunk.usage.input_tokens_details, "cached_tokens"
-                        ):
-                            usage_stats["cache_read_input_tokens"] = (
-                                chunk.usage.input_tokens_details.cached_tokens
-                            )
+                    if content is not None:
+                        final_content.append(content)
 
                     yield chunk
 
@@ -159,6 +154,7 @@ class WrappedResponses:
                 end_time = time.time()
                 latency = end_time - start_time
                 output = final_content
+
                 await self._capture_streaming_event(
                     posthog_distinct_id,
                     posthog_trace_id,
@@ -182,7 +178,7 @@ class WrappedResponses:
         posthog_privacy_mode: bool,
         posthog_groups: Optional[Dict[str, Any]],
         kwargs: Dict[str, Any],
-        usage_stats: Dict[str, int],
+        usage_stats: TokenUsage,
         latency: float,
         output: Any,
         available_tool_calls: Optional[List[Dict[str, Any]]] = None,
@@ -195,12 +191,14 @@ class WrappedResponses:
             "$ai_model": kwargs.get("model"),
             "$ai_model_parameters": get_model_params(kwargs),
             "$ai_input": with_privacy_mode(
-                self._client._ph_client, posthog_privacy_mode, kwargs.get("input")
+                self._client._ph_client,
+                posthog_privacy_mode,
+                sanitize_openai_response(kwargs.get("input")),
             ),
             "$ai_output_choices": with_privacy_mode(
                 self._client._ph_client,
                 posthog_privacy_mode,
-                output,
+                format_openai_streaming_output(output, "responses"),
             ),
             "$ai_http_status": 200,
             "$ai_input_tokens": usage_stats.get("input_tokens", 0),
@@ -340,8 +338,9 @@ class WrappedCompletions:
         **kwargs: Any,
     ):
         start_time = time.time()
-        usage_stats: Dict[str, int] = {}
+        usage_stats: TokenUsage = TokenUsage()
         accumulated_content = []
+        accumulated_tool_calls: Dict[int, Dict[str, Any]] = {}
 
         if "stream_options" not in kwargs:
             kwargs["stream_options"] = {}
@@ -351,50 +350,40 @@ class WrappedCompletions:
         async def async_generator():
             nonlocal usage_stats
             nonlocal accumulated_content  # noqa: F824
+            nonlocal accumulated_tool_calls
 
             try:
                 async for chunk in response:
-                    if hasattr(chunk, "usage") and chunk.usage:
-                        usage_stats = {
-                            k: getattr(chunk.usage, k, 0)
-                            for k in [
-                                "prompt_tokens",
-                                "completion_tokens",
-                                "total_tokens",
-                            ]
-                        }
+                    # Extract usage stats from chunk
+                    chunk_usage = extract_openai_usage_from_chunk(chunk, "chat")
+                    if chunk_usage:
+                        merge_usage_stats(usage_stats, chunk_usage)
 
-                        # Add support for cached tokens
-                        if hasattr(chunk.usage, "prompt_tokens_details") and hasattr(
-                            chunk.usage.prompt_tokens_details, "cached_tokens"
-                        ):
-                            usage_stats["cache_read_input_tokens"] = (
-                                chunk.usage.prompt_tokens_details.cached_tokens
-                            )
+                    # Extract content from chunk
+                    content = extract_openai_content_from_chunk(chunk, "chat")
+                    if content is not None:
+                        accumulated_content.append(content)
 
-                        if hasattr(chunk.usage, "output_tokens_details") and hasattr(
-                            chunk.usage.output_tokens_details, "reasoning_tokens"
-                        ):
-                            usage_stats["reasoning_tokens"] = (
-                                chunk.usage.output_tokens_details.reasoning_tokens
-                            )
-
-                    if (
-                        hasattr(chunk, "choices")
-                        and chunk.choices
-                        and len(chunk.choices) > 0
-                    ):
-                        if chunk.choices[0].delta and chunk.choices[0].delta.content:
-                            content = chunk.choices[0].delta.content
-                            if content:
-                                accumulated_content.append(content)
+                    # Extract and accumulate tool calls from chunk
+                    chunk_tool_calls = extract_openai_tool_calls_from_chunk(chunk)
+                    if chunk_tool_calls:
+                        accumulate_openai_tool_calls(
+                            accumulated_tool_calls, chunk_tool_calls
+                        )
 
                     yield chunk
 
             finally:
                 end_time = time.time()
                 latency = end_time - start_time
-                output = "".join(accumulated_content)
+
+                # Convert accumulated tool calls dict to list
+                tool_calls_list = (
+                    list(accumulated_tool_calls.values())
+                    if accumulated_tool_calls
+                    else None
+                )
+
                 await self._capture_streaming_event(
                     posthog_distinct_id,
                     posthog_trace_id,
@@ -404,7 +393,8 @@ class WrappedCompletions:
                     kwargs,
                     usage_stats,
                     latency,
-                    output,
+                    accumulated_content,
+                    tool_calls_list,
                     extract_available_tool_calls("openai", kwargs),
                 )
 
@@ -418,9 +408,10 @@ class WrappedCompletions:
         posthog_privacy_mode: bool,
         posthog_groups: Optional[Dict[str, Any]],
         kwargs: Dict[str, Any],
-        usage_stats: Dict[str, int],
+        usage_stats: TokenUsage,
         latency: float,
         output: Any,
+        tool_calls: Optional[List[Dict[str, Any]]] = None,
         available_tool_calls: Optional[List[Dict[str, Any]]] = None,
     ):
         if posthog_trace_id is None:
@@ -431,16 +422,18 @@ class WrappedCompletions:
             "$ai_model": kwargs.get("model"),
             "$ai_model_parameters": get_model_params(kwargs),
             "$ai_input": with_privacy_mode(
-                self._client._ph_client, posthog_privacy_mode, kwargs.get("messages")
+                self._client._ph_client,
+                posthog_privacy_mode,
+                sanitize_openai(kwargs.get("messages")),
             ),
             "$ai_output_choices": with_privacy_mode(
                 self._client._ph_client,
                 posthog_privacy_mode,
-                [{"content": output, "role": "assistant"}],
+                format_openai_streaming_output(output, "chat", tool_calls),
             ),
             "$ai_http_status": 200,
-            "$ai_input_tokens": usage_stats.get("prompt_tokens", 0),
-            "$ai_output_tokens": usage_stats.get("completion_tokens", 0),
+            "$ai_input_tokens": usage_stats.get("input_tokens", 0),
+            "$ai_output_tokens": usage_stats.get("output_tokens", 0),
             "$ai_cache_read_input_tokens": usage_stats.get(
                 "cache_read_input_tokens", 0
             ),
@@ -475,6 +468,7 @@ class WrappedEmbeddings:
 
     def __getattr__(self, name):
         """Fallback to original embeddings object for any methods we don't explicitly handle."""
+
         return getattr(self._original, name)
 
     async def create(
@@ -500,6 +494,7 @@ class WrappedEmbeddings:
         Returns:
             The response from OpenAI's embeddings.create call.
         """
+
         if posthog_trace_id is None:
             posthog_trace_id = str(uuid.uuid4())
 
@@ -508,12 +503,13 @@ class WrappedEmbeddings:
         end_time = time.time()
 
         # Extract usage statistics if available
-        usage_stats = {}
+        usage_stats: TokenUsage = TokenUsage()
+
         if hasattr(response, "usage") and response.usage:
-            usage_stats = {
-                "prompt_tokens": getattr(response.usage, "prompt_tokens", 0),
-                "total_tokens": getattr(response.usage, "total_tokens", 0),
-            }
+            usage_stats = TokenUsage(
+                input_tokens=getattr(response.usage, "prompt_tokens", 0),
+                output_tokens=getattr(response.usage, "completion_tokens", 0),
+            )
 
         latency = end_time - start_time
 
@@ -522,10 +518,12 @@ class WrappedEmbeddings:
             "$ai_provider": "openai",
             "$ai_model": kwargs.get("model"),
             "$ai_input": with_privacy_mode(
-                self._client._ph_client, posthog_privacy_mode, kwargs.get("input")
+                self._client._ph_client,
+                posthog_privacy_mode,
+                sanitize_openai_response(kwargs.get("input")),
             ),
             "$ai_http_status": 200,
-            "$ai_input_tokens": usage_stats.get("prompt_tokens", 0),
+            "$ai_input_tokens": usage_stats.get("input_tokens", 0),
             "$ai_latency": latency,
             "$ai_trace_id": posthog_trace_id,
             "$ai_base_url": str(self._client.base_url),
@@ -556,6 +554,7 @@ class WrappedBeta:
 
     def __getattr__(self, name):
         """Fallback to original beta object for any methods we don't explicitly handle."""
+
         return getattr(self._original, name)
 
     @property
@@ -572,6 +571,7 @@ class WrappedBetaChat:
 
     def __getattr__(self, name):
         """Fallback to original beta chat object for any methods we don't explicitly handle."""
+
         return getattr(self._original, name)
 
     @property
@@ -588,6 +588,7 @@ class WrappedBetaCompletions:
 
     def __getattr__(self, name):
         """Fallback to original beta completions object for any methods we don't explicitly handle."""
+
         return getattr(self._original, name)
 
     async def parse(

posthog/ai/openai/openai_converter.py

@@ -0,0 +1,611 @@
+"""
+OpenAI-specific conversion utilities.
+
+This module handles the conversion of OpenAI API responses and inputs
+into standardized formats for PostHog tracking. It supports both
+Chat Completions API and Responses API formats.
+"""
+
+from typing import Any, Dict, List, Optional
+
+from posthog.ai.types import (
+    FormattedContentItem,
+    FormattedFunctionCall,
+    FormattedImageContent,
+    FormattedMessage,
+    FormattedTextContent,
+    TokenUsage,
+)
+
+
+def format_openai_response(response: Any) -> List[FormattedMessage]:
+    """
+    Format an OpenAI response into standardized message format.
+
+    Handles both Chat Completions API and Responses API formats.
+
+    Args:
+        response: The response object from OpenAI API
+
+    Returns:
+        List of formatted messages with role and content
+    """
+
+    output: List[FormattedMessage] = []
+
+    if response is None:
+        return output
+
+    # Handle Chat Completions response format
+    if hasattr(response, "choices"):
+        content: List[FormattedContentItem] = []
+        role = "assistant"
+
+        for choice in response.choices:
+            if hasattr(choice, "message") and choice.message:
+                if choice.message.role:
+                    role = choice.message.role
+
+                if choice.message.content:
+                    content.append(
+                        {
+                            "type": "text",
+                            "text": choice.message.content,
+                        }
+                    )
+
+                if hasattr(choice.message, "tool_calls") and choice.message.tool_calls:
+                    for tool_call in choice.message.tool_calls:
+                        content.append(
+                            {
+                                "type": "function",
+                                "id": tool_call.id,
+                                "function": {
+                                    "name": tool_call.function.name,
+                                    "arguments": tool_call.function.arguments,
+                                },
+                            }
+                        )
+
+        if content:
+            output.append(
+                {
+                    "role": role,
+                    "content": content,
+                }
+            )
+
+    # Handle Responses API format
+    if hasattr(response, "output"):
+        content = []
+        role = "assistant"
+
+        for item in response.output:
+            if item.type == "message":
+                role = item.role
+
+                if hasattr(item, "content") and isinstance(item.content, list):
+                    for content_item in item.content:
+                        if (
+                            hasattr(content_item, "type")
+                            and content_item.type == "output_text"
+                            and hasattr(content_item, "text")
+                        ):
+                            content.append(
+                                {
+                                    "type": "text",
+                                    "text": content_item.text,
+                                }
+                            )
+
+                        elif hasattr(content_item, "text"):
+                            content.append({"type": "text", "text": content_item.text})
+
+                        elif (
+                            hasattr(content_item, "type")
+                            and content_item.type == "input_image"
+                            and hasattr(content_item, "image_url")
+                        ):
+                            image_content: FormattedImageContent = {
+                                "type": "image",
+                                "image": content_item.image_url,
+                            }
+                            content.append(image_content)
+
+                elif hasattr(item, "content"):
+                    text_content = {"type": "text", "text": str(item.content)}
+                    content.append(text_content)
+
+            elif hasattr(item, "type") and item.type == "function_call":
+                content.append(
+                    {
+                        "type": "function",
+                        "id": getattr(item, "call_id", getattr(item, "id", "")),
+                        "function": {
+                            "name": item.name,
+                            "arguments": getattr(item, "arguments", {}),
+                        },
+                    }
+                )
+
+        if content:
+            output.append(
+                {
+                    "role": role,
+                    "content": content,
+                }
+            )
+
+    return output
+
+
+def format_openai_input(
+    messages: Optional[List[Dict[str, Any]]] = None, input_data: Optional[Any] = None
+) -> List[FormattedMessage]:
+    """
+    Format OpenAI input messages.
+
+    Handles both messages parameter (Chat Completions) and input parameter (Responses API).
+
+    Args:
+        messages: List of message dictionaries for Chat Completions API
+        input_data: Input data for Responses API
+
+    Returns:
+        List of formatted messages
+    """
+
+    formatted_messages: List[FormattedMessage] = []
+
+    # Handle Chat Completions API format
+    if messages is not None:
+        for msg in messages:
+            formatted_messages.append(
+                {
+                    "role": msg.get("role", "user"),
+                    "content": msg.get("content", ""),
+                }
+            )
+
+    # Handle Responses API format
+    if input_data is not None:
+        if isinstance(input_data, list):
+            for item in input_data:
+                role = "user"
+                content = ""
+
+                if isinstance(item, dict):
+                    role = item.get("role", "user")
+                    content = item.get("content", "")
+
+                elif isinstance(item, str):
+                    content = item
+
+                else:
+                    content = str(item)
+
+                formatted_messages.append({"role": role, "content": content})
+
+        elif isinstance(input_data, str):
+            formatted_messages.append({"role": "user", "content": input_data})
+
+        else:
+            formatted_messages.append({"role": "user", "content": str(input_data)})
+
+    return formatted_messages
+
+
+def extract_openai_tools(kwargs: Dict[str, Any]) -> Optional[Any]:
+    """
+    Extract tool definitions from OpenAI API kwargs.
+
+    Args:
+        kwargs: Keyword arguments passed to OpenAI API
+
+    Returns:
+        Tool definitions if present, None otherwise
+    """
+
+    # Check for tools parameter (newer API)
+    if "tools" in kwargs:
+        return kwargs["tools"]
+
+    # Check for functions parameter (older API)
+    if "functions" in kwargs:
+        return kwargs["functions"]
+
+    return None
+
+
+def format_openai_streaming_content(
+    accumulated_content: str, tool_calls: Optional[List[Dict[str, Any]]] = None
+) -> List[FormattedContentItem]:
+    """
+    Format content from OpenAI streaming response.
+
+    Used by streaming handlers to format accumulated content.
+
+    Args:
+        accumulated_content: Accumulated text content from streaming
+        tool_calls: Optional list of tool calls accumulated during streaming
+
+    Returns:
+        List of formatted content items
+    """
+    formatted: List[FormattedContentItem] = []
+
+    # Add text content if present
+    if accumulated_content:
+        text_content: FormattedTextContent = {
+            "type": "text",
+            "text": accumulated_content,
+        }
+        formatted.append(text_content)
+
+    # Add tool calls if present
+    if tool_calls:
+        for tool_call in tool_calls:
+            function_call: FormattedFunctionCall = {
+                "type": "function",
+                "id": tool_call.get("id"),
+                "function": tool_call.get("function", {}),
+            }
+            formatted.append(function_call)
+
+    return formatted
+
+
+def extract_openai_usage_from_response(response: Any) -> TokenUsage:
+    """
+    Extract usage statistics from a full OpenAI response (non-streaming).
+    Handles both Chat Completions and Responses API.
+
+    Args:
+        response: The complete response from OpenAI API
+
+    Returns:
+        TokenUsage with standardized usage statistics
+    """
+    if not hasattr(response, "usage"):
+        return TokenUsage(input_tokens=0, output_tokens=0)
+
+    cached_tokens = 0
+    input_tokens = 0
+    output_tokens = 0
+    reasoning_tokens = 0
+
+    # Responses API format
+    if hasattr(response.usage, "input_tokens"):
+        input_tokens = response.usage.input_tokens
+    if hasattr(response.usage, "output_tokens"):
+        output_tokens = response.usage.output_tokens
+    if hasattr(response.usage, "input_tokens_details") and hasattr(
+        response.usage.input_tokens_details, "cached_tokens"
+    ):
+        cached_tokens = response.usage.input_tokens_details.cached_tokens
+    if hasattr(response.usage, "output_tokens_details") and hasattr(
+        response.usage.output_tokens_details, "reasoning_tokens"
+    ):
+        reasoning_tokens = response.usage.output_tokens_details.reasoning_tokens
+
+    # Chat Completions format
+    if hasattr(response.usage, "prompt_tokens"):
+        input_tokens = response.usage.prompt_tokens
+    if hasattr(response.usage, "completion_tokens"):
+        output_tokens = response.usage.completion_tokens
+    if hasattr(response.usage, "prompt_tokens_details") and hasattr(
+        response.usage.prompt_tokens_details, "cached_tokens"
+    ):
+        cached_tokens = response.usage.prompt_tokens_details.cached_tokens
+    if hasattr(response.usage, "completion_tokens_details") and hasattr(
+        response.usage.completion_tokens_details, "reasoning_tokens"
+    ):
+        reasoning_tokens = response.usage.completion_tokens_details.reasoning_tokens
+
+    result = TokenUsage(
+        input_tokens=input_tokens,
+        output_tokens=output_tokens,
+    )
+
+    if cached_tokens > 0:
+        result["cache_read_input_tokens"] = cached_tokens
+    if reasoning_tokens > 0:
+        result["reasoning_tokens"] = reasoning_tokens
+
+    return result
+
+
+def extract_openai_usage_from_chunk(
+    chunk: Any, provider_type: str = "chat"
+) -> TokenUsage:
+    """
+    Extract usage statistics from an OpenAI streaming chunk.
+
+    Handles both Chat Completions and Responses API formats.
+
+    Args:
+        chunk: Streaming chunk from OpenAI API
+        provider_type: Either "chat" or "responses" to handle different API formats
+
+    Returns:
+        Dictionary of usage statistics
+    """
+
+    usage: TokenUsage = TokenUsage()
+
+    if provider_type == "chat":
+        if not hasattr(chunk, "usage") or not chunk.usage:
+            return usage
+
+        # Chat Completions API uses prompt_tokens and completion_tokens
+        # Standardize to input_tokens and output_tokens
+        usage["input_tokens"] = getattr(chunk.usage, "prompt_tokens", 0)
+        usage["output_tokens"] = getattr(chunk.usage, "completion_tokens", 0)
+
+        # Handle cached tokens
+        if hasattr(chunk.usage, "prompt_tokens_details") and hasattr(
+            chunk.usage.prompt_tokens_details, "cached_tokens"
+        ):
+            usage["cache_read_input_tokens"] = (
+                chunk.usage.prompt_tokens_details.cached_tokens
+            )
+
+        # Handle reasoning tokens
+        if hasattr(chunk.usage, "completion_tokens_details") and hasattr(
+            chunk.usage.completion_tokens_details, "reasoning_tokens"
+        ):
+            usage["reasoning_tokens"] = (
+                chunk.usage.completion_tokens_details.reasoning_tokens
+            )
+
+    elif provider_type == "responses":
+        # For Responses API, usage is only in chunk.response.usage for completed events
+        if hasattr(chunk, "type") and chunk.type == "response.completed":
+            if (
+                hasattr(chunk, "response")
+                and hasattr(chunk.response, "usage")
+                and chunk.response.usage
+            ):
+                response_usage = chunk.response.usage
+                usage["input_tokens"] = getattr(response_usage, "input_tokens", 0)
+                usage["output_tokens"] = getattr(response_usage, "output_tokens", 0)
+
+                # Handle cached tokens
+                if hasattr(response_usage, "input_tokens_details") and hasattr(
+                    response_usage.input_tokens_details, "cached_tokens"
+                ):
+                    usage["cache_read_input_tokens"] = (
+                        response_usage.input_tokens_details.cached_tokens
+                    )
+
+                # Handle reasoning tokens
+                if hasattr(response_usage, "output_tokens_details") and hasattr(
+                    response_usage.output_tokens_details, "reasoning_tokens"
+                ):
+                    usage["reasoning_tokens"] = (
+                        response_usage.output_tokens_details.reasoning_tokens
+                    )
+
+    return usage
+
+
+def extract_openai_content_from_chunk(
+    chunk: Any, provider_type: str = "chat"
+) -> Optional[str]:
+    """
+    Extract content from an OpenAI streaming chunk.
+
+    Handles both Chat Completions and Responses API formats.
+
+    Args:
+        chunk: Streaming chunk from OpenAI API
+        provider_type: Either "chat" or "responses" to handle different API formats
+
+    Returns:
+        Text content if present, None otherwise
+    """
+
+    if provider_type == "chat":
+        # Chat Completions API format
+        if (
+            hasattr(chunk, "choices")
+            and chunk.choices
+            and len(chunk.choices) > 0
+            and chunk.choices[0].delta
+            and chunk.choices[0].delta.content
+        ):
+            return chunk.choices[0].delta.content
+
+    elif provider_type == "responses":
+        # Responses API format
+        if hasattr(chunk, "type") and chunk.type == "response.completed":
+            if hasattr(chunk, "response") and chunk.response:
+                res = chunk.response
+                if res.output and len(res.output) > 0:
+                    # Return the full output for responses
+                    return res.output[0]
+
+    return None
+
+
+def extract_openai_tool_calls_from_chunk(chunk: Any) -> Optional[List[Dict[str, Any]]]:
+    """
+    Extract tool calls from an OpenAI streaming chunk.
+
+    Args:
+        chunk: Streaming chunk from OpenAI API
+
+    Returns:
+        List of tool call deltas if present, None otherwise
+    """
+    if (
+        hasattr(chunk, "choices")
+        and chunk.choices
+        and len(chunk.choices) > 0
+        and chunk.choices[0].delta
+        and hasattr(chunk.choices[0].delta, "tool_calls")
+        and chunk.choices[0].delta.tool_calls
+    ):
+        tool_calls = []
+        for tool_call in chunk.choices[0].delta.tool_calls:
+            tc_dict = {
+                "index": getattr(tool_call, "index", None),
+            }
+
+            if hasattr(tool_call, "id") and tool_call.id:
+                tc_dict["id"] = tool_call.id
+
+            if hasattr(tool_call, "type") and tool_call.type:
+                tc_dict["type"] = tool_call.type
+
+            if hasattr(tool_call, "function") and tool_call.function:
+                function_dict = {}
+                if hasattr(tool_call.function, "name") and tool_call.function.name:
+                    function_dict["name"] = tool_call.function.name
+                if (
+                    hasattr(tool_call.function, "arguments")
+                    and tool_call.function.arguments
+                ):
+                    function_dict["arguments"] = tool_call.function.arguments
+                tc_dict["function"] = function_dict
+
+            tool_calls.append(tc_dict)
+        return tool_calls
+
+    return None
+
+
+def accumulate_openai_tool_calls(
+    accumulated_tool_calls: Dict[int, Dict[str, Any]],
+    chunk_tool_calls: List[Dict[str, Any]],
+) -> None:
+    """
+    Accumulate tool calls from streaming chunks.
+
+    OpenAI sends tool calls incrementally:
+    - First chunk has id, type, function.name and partial function.arguments
+    - Subsequent chunks have more function.arguments
+
+    Args:
+        accumulated_tool_calls: Dictionary mapping index to accumulated tool call data
+        chunk_tool_calls: List of tool call deltas from current chunk
+    """
+    for tool_call_delta in chunk_tool_calls:
+        index = tool_call_delta.get("index")
+        if index is None:
+            continue
+
+        # Initialize tool call if first time seeing this index
+        if index not in accumulated_tool_calls:
+            accumulated_tool_calls[index] = {
+                "id": "",
+                "type": "function",
+                "function": {
+                    "name": "",
+                    "arguments": "",
+                },
+            }
+
+        # Update with new data from delta
+        tc = accumulated_tool_calls[index]
+
+        if "id" in tool_call_delta and tool_call_delta["id"]:
+            tc["id"] = tool_call_delta["id"]
+
+        if "type" in tool_call_delta and tool_call_delta["type"]:
+            tc["type"] = tool_call_delta["type"]
+
+        if "function" in tool_call_delta:
+            func_delta = tool_call_delta["function"]
+            if "name" in func_delta and func_delta["name"]:
+                tc["function"]["name"] = func_delta["name"]
+            if "arguments" in func_delta and func_delta["arguments"]:
+                # Arguments are sent incrementally, concatenate them
+                tc["function"]["arguments"] += func_delta["arguments"]
+
+
+def format_openai_streaming_output(
+    accumulated_content: Any,
+    provider_type: str = "chat",
+    tool_calls: Optional[List[Dict[str, Any]]] = None,
+) -> List[FormattedMessage]:
+    """
+    Format the final output from OpenAI streaming.
+
+    Args:
+        accumulated_content: Accumulated content from streaming (string for chat, list for responses)
+        provider_type: Either "chat" or "responses" to handle different API formats
+        tool_calls: Optional list of accumulated tool calls
+
+    Returns:
+        List of formatted messages
+    """
+
+    if provider_type == "chat":
+        content_items: List[FormattedContentItem] = []
+
+        # Add text content if present
+        if isinstance(accumulated_content, str) and accumulated_content:
+            content_items.append({"type": "text", "text": accumulated_content})
+        elif isinstance(accumulated_content, list):
+            # If it's a list of strings, join them
+            text = "".join(str(item) for item in accumulated_content if item)
+            if text:
+                content_items.append({"type": "text", "text": text})
+
+        # Add tool calls if present
+        if tool_calls:
+            for tool_call in tool_calls:
+                if "function" in tool_call:
+                    function_call: FormattedFunctionCall = {
+                        "type": "function",
+                        "id": tool_call.get("id", ""),
+                        "function": tool_call["function"],
+                    }
+                    content_items.append(function_call)
+
+        # Return formatted message with content
+        if content_items:
+            return [{"role": "assistant", "content": content_items}]
+        else:
+            # Empty response
+            return [{"role": "assistant", "content": []}]
+
+    elif provider_type == "responses":
+        # Responses API: accumulated_content is a list of output items
+        if isinstance(accumulated_content, list) and accumulated_content:
+            # The output is already formatted, just return it
+            return accumulated_content
+        elif isinstance(accumulated_content, str):
+            return [
+                {
+                    "role": "assistant",
+                    "content": [{"type": "text", "text": accumulated_content}],
+                }
+            ]
+
+    # Fallback for any other format
+    return [
+        {
+            "role": "assistant",
+            "content": [{"type": "text", "text": str(accumulated_content)}],
+        }
+    ]
+
+
+def format_openai_streaming_input(
+    kwargs: Dict[str, Any], api_type: str = "chat"
+) -> Any:
+    """
+    Format OpenAI streaming input based on API type.
+
+    Args:
+        kwargs: Keyword arguments passed to OpenAI API
+        api_type: Either "chat" or "responses"
+
+    Returns:
+        Formatted input ready for PostHog tracking
+    """
+    from posthog.ai.utils import merge_system_prompt
+
+    return merge_system_prompt(kwargs, "openai")

posthog/ai/sanitization.py

@@ -0,0 +1,226 @@
+import re
+from typing import Any
+from urllib.parse import urlparse
+
+REDACTED_IMAGE_PLACEHOLDER = "[base64 image redacted]"
+
+
+def is_base64_data_url(text: str) -> bool:
+    return re.match(r"^data:([^;]+);base64,", text) is not None
+
+
+def is_valid_url(text: str) -> bool:
+    try:
+        result = urlparse(text)
+        return bool(result.scheme and result.netloc)
+    except Exception:
+        pass
+
+    return text.startswith(("/", "./", "../"))
+
+
+def is_raw_base64(text: str) -> bool:
+    if is_valid_url(text):
+        return False
+
+    return len(text) > 20 and re.match(r"^[A-Za-z0-9+/]+=*$", text) is not None
+
+
+def redact_base64_data_url(value: Any) -> Any:
+    if not isinstance(value, str):
+        return value
+
+    if is_base64_data_url(value):
+        return REDACTED_IMAGE_PLACEHOLDER
+
+    if is_raw_base64(value):
+        return REDACTED_IMAGE_PLACEHOLDER
+
+    return value
+
+
+def process_messages(messages: Any, transform_content_func) -> Any:
+    if not messages:
+        return messages
+
+    def process_content(content: Any) -> Any:
+        if isinstance(content, str):
+            return content
+
+        if not content:
+            return content
+
+        if isinstance(content, list):
+            return [transform_content_func(item) for item in content]
+
+        return transform_content_func(content)
+
+    def process_message(msg: Any) -> Any:
+        if not isinstance(msg, dict) or "content" not in msg:
+            return msg
+        return {**msg, "content": process_content(msg["content"])}
+
+    if isinstance(messages, list):
+        return [process_message(msg) for msg in messages]
+
+    return process_message(messages)
+
+
+def sanitize_openai_image(item: Any) -> Any:
+    if not isinstance(item, dict):
+        return item
+
+    if (
+        item.get("type") == "image_url"
+        and isinstance(item.get("image_url"), dict)
+        and "url" in item["image_url"]
+    ):
+        return {
+            **item,
+            "image_url": {
+                **item["image_url"],
+                "url": redact_base64_data_url(item["image_url"]["url"]),
+            },
+        }
+
+    return item
+
+
+def sanitize_openai_response_image(item: Any) -> Any:
+    if not isinstance(item, dict):
+        return item
+
+    if item.get("type") == "input_image" and "image_url" in item:
+        return {
+            **item,
+            "image_url": redact_base64_data_url(item["image_url"]),
+        }
+
+    return item
+
+
+def sanitize_anthropic_image(item: Any) -> Any:
+    if not isinstance(item, dict):
+        return item
+
+    if (
+        item.get("type") == "image"
+        and isinstance(item.get("source"), dict)
+        and item["source"].get("type") == "base64"
+        and "data" in item["source"]
+    ):
+        # For Anthropic, if the source type is "base64", we should always redact the data
+        # The provider is explicitly telling us this is base64 data
+        return {
+            **item,
+            "source": {
+                **item["source"],
+                "data": REDACTED_IMAGE_PLACEHOLDER,
+            },
+        }
+
+    return item
+
+
+def sanitize_gemini_part(part: Any) -> Any:
+    if not isinstance(part, dict):
+        return part
+
+    if (
+        "inline_data" in part
+        and isinstance(part["inline_data"], dict)
+        and "data" in part["inline_data"]
+    ):
+        # For Gemini, the inline_data structure indicates base64 data
+        # We should redact any string data in this context
+        return {
+            **part,
+            "inline_data": {
+                **part["inline_data"],
+                "data": REDACTED_IMAGE_PLACEHOLDER,
+            },
+        }
+
+    return part
+
+
+def process_gemini_item(item: Any) -> Any:
+    if not isinstance(item, dict):
+        return item
+
+    if "parts" in item and item["parts"]:
+        parts = item["parts"]
+        if isinstance(parts, list):
+            parts = [sanitize_gemini_part(part) for part in parts]
+        else:
+            parts = sanitize_gemini_part(parts)
+
+        return {**item, "parts": parts}
+
+    return item
+
+
+def sanitize_langchain_image(item: Any) -> Any:
+    if not isinstance(item, dict):
+        return item
+
+    if (
+        item.get("type") == "image_url"
+        and isinstance(item.get("image_url"), dict)
+        and "url" in item["image_url"]
+    ):
+        return {
+            **item,
+            "image_url": {
+                **item["image_url"],
+                "url": redact_base64_data_url(item["image_url"]["url"]),
+            },
+        }
+
+    if item.get("type") == "image" and "data" in item:
+        return {**item, "data": redact_base64_data_url(item["data"])}
+
+    if (
+        item.get("type") == "image"
+        and isinstance(item.get("source"), dict)
+        and "data" in item["source"]
+    ):
+        # Anthropic style - raw base64 in structured format, always redact
+        return {
+            **item,
+            "source": {
+                **item["source"],
+                "data": REDACTED_IMAGE_PLACEHOLDER,
+            },
+        }
+
+    if item.get("type") == "media" and "data" in item:
+        return {**item, "data": redact_base64_data_url(item["data"])}
+
+    return item
+
+
+def sanitize_openai(data: Any) -> Any:
+    return process_messages(data, sanitize_openai_image)
+
+
+def sanitize_openai_response(data: Any) -> Any:
+    return process_messages(data, sanitize_openai_response_image)
+
+
+def sanitize_anthropic(data: Any) -> Any:
+    return process_messages(data, sanitize_anthropic_image)
+
+
+def sanitize_gemini(data: Any) -> Any:
+    if not data:
+        return data
+
+    if isinstance(data, list):
+        return [process_gemini_item(item) for item in data]
+
+    return process_gemini_item(data)
+
+
+def sanitize_langchain(data: Any) -> Any:
+    return process_messages(data, sanitize_langchain_image)

posthog/ai/types.py

@@ -0,0 +1,124 @@
+"""
+Common type definitions for PostHog AI SDK.
+
+These types are used for formatting messages and responses across different AI providers
+(Anthropic, OpenAI, Gemini, etc.) to ensure consistency in tracking and data structure.
+"""
+
+from typing import Any, Dict, List, Optional, TypedDict, Union
+
+
+class FormattedTextContent(TypedDict):
+    """Formatted text content item."""
+
+    type: str  # Literal["text"]
+    text: str
+
+
+class FormattedFunctionCall(TypedDict, total=False):
+    """Formatted function/tool call content item."""
+
+    type: str  # Literal["function"]
+    id: Optional[str]
+    function: Dict[str, Any]  # Contains 'name' and 'arguments'
+
+
+class FormattedImageContent(TypedDict):
+    """Formatted image content item."""
+
+    type: str  # Literal["image"]
+    image: str
+
+
+# Union type for all formatted content items
+FormattedContentItem = Union[
+    FormattedTextContent,
+    FormattedFunctionCall,
+    FormattedImageContent,
+    Dict[str, Any],  # Fallback for unknown content types
+]
+
+
+class FormattedMessage(TypedDict):
+    """
+    Standardized message format for PostHog tracking.
+
+    Used across all providers to ensure consistent message structure
+    when sending events to PostHog.
+    """
+
+    role: str
+    content: Union[str, List[FormattedContentItem], Any]
+
+
+class TokenUsage(TypedDict, total=False):
+    """
+    Token usage information for AI model responses.
+
+    Different providers may populate different fields.
+    """
+
+    input_tokens: int
+    output_tokens: int
+    cache_read_input_tokens: Optional[int]
+    cache_creation_input_tokens: Optional[int]
+    reasoning_tokens: Optional[int]
+
+
+class ProviderResponse(TypedDict, total=False):
+    """
+    Standardized provider response format.
+
+    Used for consistent response formatting across all providers.
+    """
+
+    messages: List[FormattedMessage]
+    usage: TokenUsage
+    error: Optional[str]
+
+
+class StreamingContentBlock(TypedDict, total=False):
+    """
+    Content block used during streaming to accumulate content.
+
+    Used for tracking text and function calls as they stream in.
+    """
+
+    type: str  # "text" or "function"
+    text: Optional[str]
+    id: Optional[str]
+    function: Optional[Dict[str, Any]]
+
+
+class ToolInProgress(TypedDict):
+    """
+    Tracks a tool/function call being accumulated during streaming.
+
+    Used by Anthropic to accumulate JSON input for tools.
+    """
+
+    block: StreamingContentBlock
+    input_string: str
+
+
+class StreamingEventData(TypedDict):
+    """
+    Standardized data for streaming events across all providers.
+
+    This type ensures consistent data structure when capturing streaming events,
+    with all provider-specific formatting already completed.
+    """
+
+    provider: str  # "openai", "anthropic", "gemini"
+    model: str
+    base_url: str
+    kwargs: Dict[str, Any]  # Original kwargs for tool extraction and special handling
+    formatted_input: Any  # Provider-formatted input ready for tracking
+    formatted_output: Any  # Provider-formatted output ready for tracking
+    usage_stats: TokenUsage
+    latency: float
+    distinct_id: Optional[str]
+    trace_id: Optional[str]
+    properties: Optional[Dict[str, Any]]
+    privacy_mode: bool
+    groups: Optional[Dict[str, Any]]

posthog/ai/utils.py

@@ -1,10 +1,74 @@
 import time
 import uuid
-from typing import Any, Callable, Dict, List, Optional
-
-from httpx import URL
+from typing import Any, Callable, Dict, List, Optional, cast
 
 from posthog.client import Client as PostHogClient
+from posthog.ai.types import FormattedMessage, StreamingEventData, TokenUsage
+from posthog.ai.sanitization import (
+    sanitize_openai,
+    sanitize_anthropic,
+    sanitize_gemini,
+    sanitize_langchain,
+)
+
+
+def merge_usage_stats(
+    target: TokenUsage, source: TokenUsage, mode: str = "incremental"
+) -> None:
+    """
+    Merge streaming usage statistics into target dict, handling None values.
+
+    Supports two modes:
+    - "incremental": Add source values to target (for APIs that report new tokens)
+    - "cumulative": Replace target with source values (for APIs that report totals)
+
+    Args:
+        target: Dictionary to update with usage stats
+        source: TokenUsage that may contain None values
+        mode: Either "incremental" or "cumulative"
+    """
+    if mode == "incremental":
+        # Add new values to existing totals
+        source_input = source.get("input_tokens")
+        if source_input is not None:
+            current = target.get("input_tokens") or 0
+            target["input_tokens"] = current + source_input
+
+        source_output = source.get("output_tokens")
+        if source_output is not None:
+            current = target.get("output_tokens") or 0
+            target["output_tokens"] = current + source_output
+
+        source_cache_read = source.get("cache_read_input_tokens")
+        if source_cache_read is not None:
+            current = target.get("cache_read_input_tokens") or 0
+            target["cache_read_input_tokens"] = current + source_cache_read
+
+        source_cache_creation = source.get("cache_creation_input_tokens")
+        if source_cache_creation is not None:
+            current = target.get("cache_creation_input_tokens") or 0
+            target["cache_creation_input_tokens"] = current + source_cache_creation
+
+        source_reasoning = source.get("reasoning_tokens")
+        if source_reasoning is not None:
+            current = target.get("reasoning_tokens") or 0
+            target["reasoning_tokens"] = current + source_reasoning
+    elif mode == "cumulative":
+        # Replace with latest values (already cumulative)
+        if source.get("input_tokens") is not None:
+            target["input_tokens"] = source["input_tokens"]
+        if source.get("output_tokens") is not None:
+            target["output_tokens"] = source["output_tokens"]
+        if source.get("cache_read_input_tokens") is not None:
+            target["cache_read_input_tokens"] = source["cache_read_input_tokens"]
+        if source.get("cache_creation_input_tokens") is not None:
+            target["cache_creation_input_tokens"] = source[
+                "cache_creation_input_tokens"
+            ]
+        if source.get("reasoning_tokens") is not None:
+            target["reasoning_tokens"] = source["reasoning_tokens"]
+    else:
+        raise ValueError(f"Invalid mode: {mode}. Must be 'incremental' or 'cumulative'")
 
 
 def get_model_params(kwargs: Dict[str, Any]) -> Dict[str, Any]:
@@ -29,349 +93,135 @@ def get_model_params(kwargs: Dict[str, Any]) -> Dict[str, Any]:
     return model_params
 
 
-def get_usage(response, provider: str) -> Dict[str, Any]:
+def get_usage(response, provider: str) -> TokenUsage:
+    """
+    Extract usage statistics from response based on provider.
+    Delegates to provider-specific converter functions.
+    """
     if provider == "anthropic":
-        return {
-            "input_tokens": response.usage.input_tokens,
-            "output_tokens": response.usage.output_tokens,
-            "cache_read_input_tokens": response.usage.cache_read_input_tokens,
-            "cache_creation_input_tokens": response.usage.cache_creation_input_tokens,
-        }
+        from posthog.ai.anthropic.anthropic_converter import (
+            extract_anthropic_usage_from_response,
+        )
+
+        return extract_anthropic_usage_from_response(response)
     elif provider == "openai":
-        cached_tokens = 0
-        input_tokens = 0
-        output_tokens = 0
-        reasoning_tokens = 0
+        from posthog.ai.openai.openai_converter import (
+            extract_openai_usage_from_response,
+        )
 
-        # responses api
-        if hasattr(response.usage, "input_tokens"):
-            input_tokens = response.usage.input_tokens
-        if hasattr(response.usage, "output_tokens"):
-            output_tokens = response.usage.output_tokens
-        if hasattr(response.usage, "input_tokens_details") and hasattr(
-            response.usage.input_tokens_details, "cached_tokens"
-        ):
-            cached_tokens = response.usage.input_tokens_details.cached_tokens
-        if hasattr(response.usage, "output_tokens_details") and hasattr(
-            response.usage.output_tokens_details, "reasoning_tokens"
-        ):
-            reasoning_tokens = response.usage.output_tokens_details.reasoning_tokens
-
-        # chat completions
-        if hasattr(response.usage, "prompt_tokens"):
-            input_tokens = response.usage.prompt_tokens
-        if hasattr(response.usage, "completion_tokens"):
-            output_tokens = response.usage.completion_tokens
-        if hasattr(response.usage, "prompt_tokens_details") and hasattr(
-            response.usage.prompt_tokens_details, "cached_tokens"
-        ):
-            cached_tokens = response.usage.prompt_tokens_details.cached_tokens
-
-        return {
-            "input_tokens": input_tokens,
-            "output_tokens": output_tokens,
-            "cache_read_input_tokens": cached_tokens,
-            "reasoning_tokens": reasoning_tokens,
-        }
+        return extract_openai_usage_from_response(response)
     elif provider == "gemini":
-        input_tokens = 0
-        output_tokens = 0
+        from posthog.ai.gemini.gemini_converter import (
+            extract_gemini_usage_from_response,
+        )
 
-        if hasattr(response, "usage_metadata") and response.usage_metadata:
-            input_tokens = getattr(response.usage_metadata, "prompt_token_count", 0)
-            output_tokens = getattr(
-                response.usage_metadata, "candidates_token_count", 0
-            )
+        return extract_gemini_usage_from_response(response)
 
-        return {
-            "input_tokens": input_tokens,
-            "output_tokens": output_tokens,
-            "cache_read_input_tokens": 0,
-            "cache_creation_input_tokens": 0,
-            "reasoning_tokens": 0,
-        }
-    return {
-        "input_tokens": 0,
-        "output_tokens": 0,
-        "cache_read_input_tokens": 0,
-        "cache_creation_input_tokens": 0,
-        "reasoning_tokens": 0,
-    }
+    return TokenUsage(input_tokens=0, output_tokens=0)
 
 
 def format_response(response, provider: str):
     """
     Format a regular (non-streaming) response.
     """
-    output = []
-    if response is None:
-        return output
     if provider == "anthropic":
-        return format_response_anthropic(response)
+        from posthog.ai.anthropic.anthropic_converter import format_anthropic_response
+
+        return format_anthropic_response(response)
     elif provider == "openai":
-        return format_response_openai(response)
+        from posthog.ai.openai.openai_converter import format_openai_response
+
+        return format_openai_response(response)
     elif provider == "gemini":
-        return format_response_gemini(response)
-    return output
+        from posthog.ai.gemini.gemini_converter import format_gemini_response
 
-
-def format_response_anthropic(response):
-    output = []
-    content = []
-
-    for choice in response.content:
-        if (
-            hasattr(choice, "type")
-            and choice.type == "text"
-            and hasattr(choice, "text")
-            and choice.text
-        ):
-            content.append({"type": "text", "text": choice.text})
-        elif (
-            hasattr(choice, "type")
-            and choice.type == "tool_use"
-            and hasattr(choice, "name")
-            and hasattr(choice, "id")
-        ):
-            tool_call = {
-                "type": "function",
-                "id": choice.id,
-                "function": {
-                    "name": choice.name,
-                    "arguments": getattr(choice, "input", {}),
-                },
-            }
-            content.append(tool_call)
-
-    if content:
-        message = {
-            "role": "assistant",
-            "content": content,
-        }
-        output.append(message)
-
-    return output
-
-
-def format_response_openai(response):
-    output = []
-
-    if hasattr(response, "choices"):
-        content = []
-        role = "assistant"
-
-        for choice in response.choices:
-            # Handle Chat Completions response format
-            if hasattr(choice, "message") and choice.message:
-                if choice.message.role:
-                    role = choice.message.role
-
-                if choice.message.content:
-                    content.append({"type": "text", "text": choice.message.content})
-
-                if hasattr(choice.message, "tool_calls") and choice.message.tool_calls:
-                    for tool_call in choice.message.tool_calls:
-                        content.append(
-                            {
-                                "type": "function",
-                                "id": tool_call.id,
-                                "function": {
-                                    "name": tool_call.function.name,
-                                    "arguments": tool_call.function.arguments,
-                                },
-                            }
-                        )
-
-        if content:
-            message = {
-                "role": role,
-                "content": content,
-            }
-            output.append(message)
-
-    # Handle Responses API format
-    if hasattr(response, "output"):
-        content = []
-        role = "assistant"
-
-        for item in response.output:
-            if item.type == "message":
-                role = item.role
-
-                if hasattr(item, "content") and isinstance(item.content, list):
-                    for content_item in item.content:
-                        if (
-                            hasattr(content_item, "type")
-                            and content_item.type == "output_text"
-                            and hasattr(content_item, "text")
-                        ):
-                            content.append({"type": "text", "text": content_item.text})
-                        elif hasattr(content_item, "text"):
-                            content.append({"type": "text", "text": content_item.text})
-                        elif (
-                            hasattr(content_item, "type")
-                            and content_item.type == "input_image"
-                            and hasattr(content_item, "image_url")
-                        ):
-                            content.append(
-                                {
-                                    "type": "image",
-                                    "image": content_item.image_url,
-                                }
-                            )
-                elif hasattr(item, "content"):
-                    content.append({"type": "text", "text": str(item.content)})
-
-            elif hasattr(item, "type") and item.type == "function_call":
-                content.append(
-                    {
-                        "type": "function",
-                        "id": getattr(item, "call_id", getattr(item, "id", "")),
-                        "function": {
-                            "name": item.name,
-                            "arguments": getattr(item, "arguments", {}),
-                        },
-                    }
-                )
-
-        if content:
-            message = {
-                "role": role,
-                "content": content,
-            }
-            output.append(message)
-
-    return output
-
-
-def format_response_gemini(response):
-    output = []
-
-    if hasattr(response, "candidates") and response.candidates:
-        for candidate in response.candidates:
-            if hasattr(candidate, "content") and candidate.content:
-                content = []
-
-                if hasattr(candidate.content, "parts") and candidate.content.parts:
-                    for part in candidate.content.parts:
-                        if hasattr(part, "text") and part.text:
-                            content.append({"type": "text", "text": part.text})
-                        elif hasattr(part, "function_call") and part.function_call:
-                            function_call = part.function_call
-                            content.append(
-                                {
-                                    "type": "function",
-                                    "function": {
-                                        "name": function_call.name,
-                                        "arguments": function_call.args,
-                                    },
-                                }
-                            )
-
-                if content:
-                    message = {
-                        "role": "assistant",
-                        "content": content,
-                    }
-                    output.append(message)
-
-            elif hasattr(candidate, "text") and candidate.text:
-                output.append(
-                    {
-                        "role": "assistant",
-                        "content": [{"type": "text", "text": candidate.text}],
-                    }
-                )
-    elif hasattr(response, "text") and response.text:
-        output.append(
-            {
-                "role": "assistant",
-                "content": [{"type": "text", "text": response.text}],
-            }
-        )
-
-    return output
+        return format_gemini_response(response)
+    return []
 
 
 def extract_available_tool_calls(provider: str, kwargs: Dict[str, Any]):
+    """
+    Extract available tool calls for the given provider.
+    """
     if provider == "anthropic":
-        if "tools" in kwargs:
-            return kwargs["tools"]
+        from posthog.ai.anthropic.anthropic_converter import extract_anthropic_tools
 
-        return None
+        return extract_anthropic_tools(kwargs)
     elif provider == "gemini":
-        if "config" in kwargs and hasattr(kwargs["config"], "tools"):
-            return kwargs["config"].tools
+        from posthog.ai.gemini.gemini_converter import extract_gemini_tools
 
-        return None
+        return extract_gemini_tools(kwargs)
     elif provider == "openai":
-        if "tools" in kwargs:
-            return kwargs["tools"]
+        from posthog.ai.openai.openai_converter import extract_openai_tools
 
-        return None
+        return extract_openai_tools(kwargs)
+    return None
 
 
-def merge_system_prompt(kwargs: Dict[str, Any], provider: str):
-    messages: List[Dict[str, Any]] = []
+def merge_system_prompt(
+    kwargs: Dict[str, Any], provider: str
+) -> List[FormattedMessage]:
+    """
+    Merge system prompts and format messages for the given provider.
+    """
     if provider == "anthropic":
+        from posthog.ai.anthropic.anthropic_converter import format_anthropic_input
+
         messages = kwargs.get("messages") or []
-        if kwargs.get("system") is None:
-            return messages
-        return [{"role": "system", "content": kwargs.get("system")}] + messages
+        system = kwargs.get("system")
+        return format_anthropic_input(messages, system)
     elif provider == "gemini":
+        from posthog.ai.gemini.gemini_converter import format_gemini_input_with_system
+
         contents = kwargs.get("contents", [])
-        if isinstance(contents, str):
-            return [{"role": "user", "content": contents}]
-        elif isinstance(contents, list):
-            formatted = []
-            for item in contents:
-                if isinstance(item, str):
-                    formatted.append({"role": "user", "content": item})
-                elif hasattr(item, "text"):
-                    formatted.append({"role": "user", "content": item.text})
-                else:
-                    formatted.append({"role": "user", "content": str(item)})
-            return formatted
-        else:
-            return [{"role": "user", "content": str(contents)}]
+        config = kwargs.get("config")
+        return format_gemini_input_with_system(contents, config)
+    elif provider == "openai":
+        from posthog.ai.openai.openai_converter import format_openai_input
 
-    # For OpenAI, handle both Chat Completions and Responses API
-    if kwargs.get("messages") is not None:
-        messages = list(kwargs.get("messages", []))
+        # For OpenAI, handle both Chat Completions and Responses API
+        messages_param = kwargs.get("messages")
+        input_param = kwargs.get("input")
 
-    if kwargs.get("input") is not None:
-        input_data = kwargs.get("input")
-        if isinstance(input_data, list):
-            messages.extend(input_data)
-        else:
-            messages.append({"role": "user", "content": input_data})
+        # Get base formatted messages
+        messages = format_openai_input(messages_param, input_param)
 
-    # Check if system prompt is provided as a separate parameter
-    if kwargs.get("system") is not None:
-        has_system = any(msg.get("role") == "system" for msg in messages)
-        if not has_system:
-            messages = [{"role": "system", "content": kwargs.get("system")}] + messages
+        # Check if system prompt is provided as a separate parameter
+        if kwargs.get("system") is not None:
+            has_system = any(msg.get("role") == "system" for msg in messages)
+            if not has_system:
+                system_msg = cast(
+                    FormattedMessage,
+                    {"role": "system", "content": kwargs.get("system")},
+                )
+                messages = [system_msg] + messages
 
-    # For Responses API, add instructions to the system prompt if provided
-    if kwargs.get("instructions") is not None:
-        # Find the system message if it exists
-        system_idx = next(
-            (i for i, msg in enumerate(messages) if msg.get("role") == "system"), None
-        )
-
-        if system_idx is not None:
-            # Append instructions to existing system message
-            system_content = messages[system_idx].get("content", "")
-            messages[system_idx]["content"] = (
-                f"{system_content}\n\n{kwargs.get('instructions')}"
+        # For Responses API, add instructions to the system prompt if provided
+        if kwargs.get("instructions") is not None:
+            # Find the system message if it exists
+            system_idx = next(
+                (i for i, msg in enumerate(messages) if msg.get("role") == "system"),
+                None,
             )
-        else:
-            # Create a new system message with instructions
-            messages = [
-                {"role": "system", "content": kwargs.get("instructions")}
-            ] + messages
 
-    return messages
+            if system_idx is not None:
+                # Append instructions to existing system message
+                system_content = messages[system_idx].get("content", "")
+                messages[system_idx]["content"] = (
+                    f"{system_content}\n\n{kwargs.get('instructions')}"
+                )
+            else:
+                # Create a new system message with instructions
+                instruction_msg = cast(
+                    FormattedMessage,
+                    {"role": "system", "content": kwargs.get("instructions")},
+                )
+                messages = [instruction_msg] + messages
+
+        return messages
+
+    # Default case - return empty list
+    return []
 
 
 def call_llm_and_track_usage(
@@ -382,7 +232,7 @@ def call_llm_and_track_usage(
     posthog_properties: Optional[Dict[str, Any]],
     posthog_privacy_mode: bool,
     posthog_groups: Optional[Dict[str, Any]],
-    base_url: URL,
+    base_url: str,
     call_method: Callable[..., Any],
     **kwargs: Any,
 ) -> Any:
@@ -394,8 +244,8 @@ def call_llm_and_track_usage(
     response = None
     error = None
     http_status = 200
-    usage: Dict[str, Any] = {}
-    error_params: Dict[str, any] = {}
+    usage: TokenUsage = TokenUsage()
+    error_params: Dict[str, Any] = {}
 
     try:
         response = call_method(**kwargs)
@@ -422,12 +272,15 @@ def call_llm_and_track_usage(
             usage = get_usage(response, provider)
 
         messages = merge_system_prompt(kwargs, provider)
+        sanitized_messages = sanitize_messages(messages, provider)
 
         event_properties = {
             "$ai_provider": provider,
             "$ai_model": kwargs.get("model"),
             "$ai_model_parameters": get_model_params(kwargs),
-            "$ai_input": with_privacy_mode(ph_client, posthog_privacy_mode, messages),
+            "$ai_input": with_privacy_mode(
+                ph_client, posthog_privacy_mode, sanitized_messages
+            ),
             "$ai_output_choices": with_privacy_mode(
                 ph_client, posthog_privacy_mode, format_response(response, provider)
             ),
@@ -446,27 +299,17 @@ def call_llm_and_track_usage(
         if available_tool_calls:
             event_properties["$ai_tools"] = available_tool_calls
 
-        if (
-            usage.get("cache_read_input_tokens") is not None
-            and usage.get("cache_read_input_tokens", 0) > 0
-        ):
-            event_properties["$ai_cache_read_input_tokens"] = usage.get(
-                "cache_read_input_tokens", 0
-            )
+        cache_read = usage.get("cache_read_input_tokens")
+        if cache_read is not None and cache_read > 0:
+            event_properties["$ai_cache_read_input_tokens"] = cache_read
 
-        if (
-            usage.get("cache_creation_input_tokens") is not None
-            and usage.get("cache_creation_input_tokens", 0) > 0
-        ):
-            event_properties["$ai_cache_creation_input_tokens"] = usage.get(
-                "cache_creation_input_tokens", 0
-            )
+        cache_creation = usage.get("cache_creation_input_tokens")
+        if cache_creation is not None and cache_creation > 0:
+            event_properties["$ai_cache_creation_input_tokens"] = cache_creation
 
-        if (
-            usage.get("reasoning_tokens") is not None
-            and usage.get("reasoning_tokens", 0) > 0
-        ):
-            event_properties["$ai_reasoning_tokens"] = usage.get("reasoning_tokens", 0)
+        reasoning = usage.get("reasoning_tokens")
+        if reasoning is not None and reasoning > 0:
+            event_properties["$ai_reasoning_tokens"] = reasoning
 
         if posthog_distinct_id is None:
             event_properties["$process_person_profile"] = False
@@ -500,7 +343,7 @@ async def call_llm_and_track_usage_async(
     posthog_properties: Optional[Dict[str, Any]],
     posthog_privacy_mode: bool,
     posthog_groups: Optional[Dict[str, Any]],
-    base_url: URL,
+    base_url: str,
     call_async_method: Callable[..., Any],
     **kwargs: Any,
 ) -> Any:
@@ -508,8 +351,8 @@ async def call_llm_and_track_usage_async(
     response = None
     error = None
     http_status = 200
-    usage: Dict[str, Any] = {}
-    error_params: Dict[str, any] = {}
+    usage: TokenUsage = TokenUsage()
+    error_params: Dict[str, Any] = {}
 
     try:
         response = await call_async_method(**kwargs)
@@ -536,12 +379,15 @@ async def call_llm_and_track_usage_async(
             usage = get_usage(response, provider)
 
         messages = merge_system_prompt(kwargs, provider)
+        sanitized_messages = sanitize_messages(messages, provider)
 
         event_properties = {
             "$ai_provider": provider,
             "$ai_model": kwargs.get("model"),
             "$ai_model_parameters": get_model_params(kwargs),
-            "$ai_input": with_privacy_mode(ph_client, posthog_privacy_mode, messages),
+            "$ai_input": with_privacy_mode(
+                ph_client, posthog_privacy_mode, sanitized_messages
+            ),
             "$ai_output_choices": with_privacy_mode(
                 ph_client, posthog_privacy_mode, format_response(response, provider)
             ),
@@ -560,21 +406,13 @@ async def call_llm_and_track_usage_async(
         if available_tool_calls:
             event_properties["$ai_tools"] = available_tool_calls
 
-        if (
-            usage.get("cache_read_input_tokens") is not None
-            and usage.get("cache_read_input_tokens", 0) > 0
-        ):
-            event_properties["$ai_cache_read_input_tokens"] = usage.get(
-                "cache_read_input_tokens", 0
-            )
+        cache_read = usage.get("cache_read_input_tokens")
+        if cache_read is not None and cache_read > 0:
+            event_properties["$ai_cache_read_input_tokens"] = cache_read
 
-        if (
-            usage.get("cache_creation_input_tokens") is not None
-            and usage.get("cache_creation_input_tokens", 0) > 0
-        ):
-            event_properties["$ai_cache_creation_input_tokens"] = usage.get(
-                "cache_creation_input_tokens", 0
-            )
+        cache_creation = usage.get("cache_creation_input_tokens")
+        if cache_creation is not None and cache_creation > 0:
+            event_properties["$ai_cache_creation_input_tokens"] = cache_creation
 
         if posthog_distinct_id is None:
             event_properties["$process_person_profile"] = False
@@ -600,7 +438,122 @@ async def call_llm_and_track_usage_async(
     return response
 
 
+def sanitize_messages(data: Any, provider: str) -> Any:
+    """Sanitize messages using provider-specific sanitization functions."""
+    if provider == "anthropic":
+        return sanitize_anthropic(data)
+    elif provider == "openai":
+        return sanitize_openai(data)
+    elif provider == "gemini":
+        return sanitize_gemini(data)
+    elif provider == "langchain":
+        return sanitize_langchain(data)
+    return data
+
+
 def with_privacy_mode(ph_client: PostHogClient, privacy_mode: bool, value: Any):
     if ph_client.privacy_mode or privacy_mode:
         return None
     return value
+
+
+def capture_streaming_event(
+    ph_client: PostHogClient,
+    event_data: StreamingEventData,
+):
+    """
+    Unified streaming event capture for all LLM providers.
+
+    This function handles the common logic for capturing streaming events across all providers.
+    All provider-specific formatting should be done BEFORE calling this function.
+
+    The function handles:
+    - Building PostHog event properties
+    - Extracting and adding tools based on provider
+    - Applying privacy mode
+    - Adding special token fields (cache, reasoning)
+    - Provider-specific fields (e.g., OpenAI instructions)
+    - Sending the event to PostHog
+
+    Args:
+        ph_client: PostHog client instance
+        event_data: Standardized streaming event data containing all necessary information
+    """
+    trace_id = event_data.get("trace_id") or str(uuid.uuid4())
+
+    # Build base event properties
+    event_properties = {
+        "$ai_provider": event_data["provider"],
+        "$ai_model": event_data["model"],
+        "$ai_model_parameters": get_model_params(event_data["kwargs"]),
+        "$ai_input": with_privacy_mode(
+            ph_client,
+            event_data["privacy_mode"],
+            event_data["formatted_input"],
+        ),
+        "$ai_output_choices": with_privacy_mode(
+            ph_client,
+            event_data["privacy_mode"],
+            event_data["formatted_output"],
+        ),
+        "$ai_http_status": 200,
+        "$ai_input_tokens": event_data["usage_stats"].get("input_tokens", 0),
+        "$ai_output_tokens": event_data["usage_stats"].get("output_tokens", 0),
+        "$ai_latency": event_data["latency"],
+        "$ai_trace_id": trace_id,
+        "$ai_base_url": str(event_data["base_url"]),
+        **(event_data.get("properties") or {}),
+    }
+
+    # Extract and add tools based on provider
+    available_tools = extract_available_tool_calls(
+        event_data["provider"],
+        event_data["kwargs"],
+    )
+    if available_tools:
+        event_properties["$ai_tools"] = available_tools
+
+    # Add optional token fields
+    # For Anthropic, always include cache fields even if 0 (backward compatibility)
+    # For others, only include if present and non-zero
+    if event_data["provider"] == "anthropic":
+        # Anthropic always includes cache fields
+        cache_read = event_data["usage_stats"].get("cache_read_input_tokens", 0)
+        cache_creation = event_data["usage_stats"].get("cache_creation_input_tokens", 0)
+        event_properties["$ai_cache_read_input_tokens"] = cache_read
+        event_properties["$ai_cache_creation_input_tokens"] = cache_creation
+    else:
+        # Other providers only include if non-zero
+        optional_token_fields = [
+            "cache_read_input_tokens",
+            "cache_creation_input_tokens",
+            "reasoning_tokens",
+        ]
+
+        for field in optional_token_fields:
+            value = event_data["usage_stats"].get(field)
+            if value is not None and isinstance(value, int) and value > 0:
+                event_properties[f"$ai_{field}"] = value
+
+    # Handle provider-specific fields
+    if (
+        event_data["provider"] == "openai"
+        and event_data["kwargs"].get("instructions") is not None
+    ):
+        event_properties["$ai_instructions"] = with_privacy_mode(
+            ph_client,
+            event_data["privacy_mode"],
+            event_data["kwargs"]["instructions"],
+        )
+
+    if event_data.get("distinct_id") is None:
+        event_properties["$process_person_profile"] = False
+
+    # Send event to PostHog
+    if hasattr(ph_client, "capture"):
+        ph_client.capture(
+            distinct_id=event_data.get("distinct_id") or trace_id,
+            event="$ai_generation",
+            properties=event_properties,
+            groups=event_data.get("groups"),
+        )

posthog/args.py

@@ -1,4 +1,4 @@
-from typing import TypedDict, Optional, Any, Dict, Union, Tuple, Type
+from typing import TypedDict, Optional, Any, Dict, List, Union, Tuple, Type
 from types import TracebackType
 from typing_extensions import NotRequired  # For Python < 3.11 compatibility
 from datetime import datetime
@@ -69,3 +69,33 @@ ExcInfo = Union[
 ]
 
 ExceptionArg = Union[BaseException, ExcInfo]
+
+
+# AI Event Types (literal strings to enforce valid event types)
+AI_EVENT_TYPE = Union[
+    str,  # Allow str for flexibility but document the expected values
+]  # "$ai_generation", "$ai_trace", "$ai_span", "$ai_embedding", "$ai_metric", "$ai_feedback"
+
+
+class OptionalCaptureAIArgs(TypedDict):
+    """Optional arguments for the capture_ai method.
+
+    Args:
+        distinct_id: Unique identifier for the person associated with this AI event. If not set, the context
+            distinct_id is used, if available, otherwise a UUID is generated.
+        properties: Dictionary of AI event properties to track. Must include required properties for the event type.
+        blob_properties: List of property names that should be sent as blobs (e.g., '$ai_input', '$ai_output_choices').
+            These properties will be extracted from `properties` and sent as multipart blobs.
+        timestamp: When the event occurred (defaults to current time)
+        uuid: Unique identifier for this specific event. If not provided, one is generated.
+        groups: Group identifiers to associate with this event (format: {group_type: group_key})
+        disable_geoip: Whether to disable GeoIP lookup for this event.
+    """
+
+    distinct_id: NotRequired[Optional[ID_TYPES]]
+    properties: NotRequired[Optional[Dict[str, Any]]]
+    blob_properties: NotRequired[Optional[List[str]]]
+    timestamp: NotRequired[Optional[Union[datetime, str]]]
+    uuid: NotRequired[Optional[str]]
+    groups: NotRequired[Optional[Dict[str, str]]]
+    disable_geoip: NotRequired[Optional[bool]]

posthog/client.py

@@ -1,16 +1,25 @@
 import atexit
+import json
 import logging
 import os
 import sys
 from datetime import datetime, timedelta
-from typing import Any, Dict, Optional, Union
+from io import BytesIO
+from typing import Any, Dict, List, Optional, Tuple, Union
 from typing_extensions import Unpack
 from uuid import uuid4
 
 from dateutil.tz import tzutc
 from six import string_types
 
-from posthog.args import OptionalCaptureArgs, OptionalSetArgs, ID_TYPES, ExceptionArg
+from posthog.args import (
+    OptionalCaptureArgs,
+    OptionalSetArgs,
+    OptionalCaptureAIArgs,
+    AI_EVENT_TYPE,
+    ID_TYPES,
+    ExceptionArg,
+)
 from posthog.consumer import Consumer
 from posthog.exception_capture import ExceptionCapture
 from posthog.exception_utils import (
@@ -20,7 +29,11 @@ from posthog.exception_utils import (
     exception_is_already_captured,
     mark_exception_as_captured,
 )
-from posthog.feature_flags import InconclusiveMatchError, match_feature_flag_properties
+from posthog.feature_flags import (
+    InconclusiveMatchError,
+    RequiresServerEvaluation,
+    match_feature_flag_properties,
+)
 from posthog.poller import Poller
 from posthog.request import (
     DEFAULT_HOST,
@@ -99,6 +112,134 @@ def add_context_tags(properties):
     return properties
 
 
+def _generate_multipart_boundary() -> str:
+    """Generate a random boundary string for multipart requests."""
+    return f"----WebKitFormBoundary{uuid4().hex[:16]}"
+
+
+def _encode_multipart_part(
+    boundary: str,
+    name: str,
+    content: bytes,
+    content_type: str,
+    filename: Optional[str] = None,
+) -> bytes:
+    """Encode a single part of a multipart/form-data request."""
+    part = f"--{boundary}\r\n".encode("utf-8")
+
+    if filename:
+        part += f'Content-Disposition: form-data; name="{name}"; filename="{filename}"\r\n'.encode(
+            "utf-8"
+        )
+    else:
+        part += f'Content-Disposition: form-data; name="{name}"\r\n'.encode("utf-8")
+
+    part += f"Content-Type: {content_type}\r\n\r\n".encode("utf-8")
+    part += content
+    part += b"\r\n"
+
+    return part
+
+
+def _build_multipart_body(
+    event_data: Dict[str, Any],
+    properties: Optional[Dict[str, Any]],
+    blob_properties: Optional[List[str]],
+) -> Tuple[bytes, str]:
+    """
+    Build a multipart/form-data body for AI capture endpoint.
+
+    Args:
+        event_data: The event data (uuid, event, distinct_id, timestamp)
+        properties: Event properties (small properties)
+        blob_properties: List of property names to send as blobs
+
+    Returns:
+        Tuple of (body_bytes, content_type_header)
+    """
+    boundary = _generate_multipart_boundary()
+    body = BytesIO()
+
+    # Part 1: Event (required, must be first)
+    event_json = json.dumps(event_data).encode("utf-8")
+    body.write(
+        _encode_multipart_part(boundary, "event", event_json, "application/json")
+    )
+
+    # Separate properties into small properties and blob properties
+    small_properties = {}
+    blob_data = {}
+
+    if properties:
+        blob_property_names = set(blob_properties or [])
+        for key, value in properties.items():
+            if key in blob_property_names:
+                blob_data[key] = value
+            else:
+                small_properties[key] = value
+
+    # Part 2: Event properties (if there are any small properties)
+    if small_properties:
+        properties_json = json.dumps(small_properties).encode("utf-8")
+        body.write(
+            _encode_multipart_part(
+                boundary, "event.properties", properties_json, "application/json"
+            )
+        )
+
+    # Part 3+: Blob parts
+    for property_name, property_value in blob_data.items():
+        # Serialize the blob data as JSON
+        blob_json = json.dumps(property_value).encode("utf-8")
+        blob_filename = f"blob_{uuid4().hex[:8]}"
+
+        body.write(
+            _encode_multipart_part(
+                boundary,
+                f"event.properties.{property_name}",
+                blob_json,
+                "application/json",
+                filename=blob_filename,
+            )
+        )
+
+    # Final boundary
+    body.write(f"--{boundary}--\r\n".encode("utf-8"))
+
+    body_bytes = body.getvalue()
+    content_type = f"multipart/form-data; boundary={boundary}"
+
+    return body_bytes, content_type
+
+
+def no_throw(default_return=None):
+    """
+    Decorator to prevent raising exceptions from public API methods.
+    Note that this doesn't prevent errors from propagating via `on_error`.
+    Exceptions will still be raised if the debug flag is enabled.
+
+    Args:
+        default_return: Value to return on exception (default: None)
+    """
+
+    def decorator(func):
+        from functools import wraps
+
+        @wraps(func)
+        def wrapper(self, *args, **kwargs):
+            try:
+                return func(self, *args, **kwargs)
+            except Exception as e:
+                if self.debug:
+                    raise e
+                self.log.exception(f"Error in {func.__name__}: {e}")
+                return default_return
+
+        return wrapper
+
+    return decorator
+
+
 class Client(object):
     """
     This is the SDK reference for the PostHog Python SDK.
@@ -481,6 +622,7 @@ class Client(object):
 
         return normalize_flags_response(resp_data)
 
+    @no_throw()
     def capture(
         self, event: str, **kwargs: Unpack[OptionalCaptureArgs]
     ) -> Optional[str]:
@@ -657,6 +799,7 @@ class Client(object):
                 f"Expected bool or dict."
             )
 
+    @no_throw()
     def set(self, **kwargs: Unpack[OptionalSetArgs]) -> Optional[str]:
         """
         Set properties on a person profile.
@@ -690,6 +833,8 @@ class Client(object):
 
         Category:
             Identification
+
+        Note: This method will not raise exceptions. Errors are logged.
         """
         distinct_id = kwargs.get("distinct_id", None)
         properties = kwargs.get("properties", None)
@@ -716,6 +861,7 @@ class Client(object):
 
         return self._enqueue(msg, disable_geoip)
 
+    @no_throw()
     def set_once(self, **kwargs: Unpack[OptionalSetArgs]) -> Optional[str]:
         """
         Set properties on a person profile only if they haven't been set before.
@@ -734,6 +880,8 @@ class Client(object):
 
         Category:
             Identification
+
+        Note: This method will not raise exceptions. Errors are logged.
         """
         distinct_id = kwargs.get("distinct_id", None)
         properties = kwargs.get("properties", None)
@@ -759,6 +907,7 @@ class Client(object):
 
         return self._enqueue(msg, disable_geoip)
 
+    @no_throw()
     def group_identify(
         self,
         group_type: str,
@@ -791,6 +940,8 @@ class Client(object):
 
         Category:
             Identification
+
+        Note: This method will not raise exceptions. Errors are logged.
         """
         properties = properties or {}
 
@@ -815,6 +966,7 @@ class Client(object):
 
         return self._enqueue(msg, disable_geoip)
 
+    @no_throw()
     def alias(
         self,
         previous_id: str,
@@ -840,6 +992,8 @@ class Client(object):
 
         Category:
             Identification
+
+        Note: This method will not raise exceptions. Errors are logged.
         """
         (distinct_id, personless) = get_identity_state(distinct_id)
 
@@ -932,7 +1086,6 @@ class Client(object):
                     "value"
                 ),
                 "$exception_list": all_exceptions_with_trace_and_in_app,
-                "$exception_personURL": f"{remove_trailing_slash(self.raw_host)}/project/{self.api_key}/person/{distinct_id}",
                 **properties,
             }
 
@@ -961,6 +1114,196 @@ class Client(object):
         except Exception as e:
             self.log.exception(f"Failed to capture exception: {e}")
 
+    @no_throw()
+    def capture_ai(
+        self, event: AI_EVENT_TYPE, **kwargs: Unpack[OptionalCaptureAIArgs]
+    ) -> Optional[str]:
+        """
+        Capture an AI event to the dedicated AI endpoint with support for large payloads.
+
+        This method sends AI events (like $ai_generation, $ai_trace, etc.) to PostHog's
+        specialized AI endpoint (/i/v0/ai) which supports large payloads through multipart/form-data
+        and blob storage in S3.
+
+        Args:
+            event: The AI event type. Must be one of: "$ai_generation", "$ai_trace", "$ai_span",
+                   "$ai_embedding", "$ai_metric", "$ai_feedback"
+            distinct_id: The distinct ID of the user.
+            properties: A dictionary of AI event properties. Must include required properties based on event type:
+                - All events: "$ai_model" (required)
+                - $ai_generation: "$ai_provider", "$ai_trace_id" (required)
+                - $ai_trace: "$ai_trace_id" (required)
+                - $ai_span: "$ai_trace_id", "$ai_span_id" (required)
+                - $ai_embedding: "$ai_provider", "$ai_trace_id" (required)
+            blob_properties: List of property names to send as blobs (large data stored in S3).
+                Common blob properties: "$ai_input", "$ai_output_choices", "$ai_input_state", "$ai_output_state"
+                If not provided, defaults to common blob properties based on event type.
+            timestamp: The timestamp of the event.
+            uuid: A unique identifier for the event.
+            groups: A dictionary of group information.
+            disable_geoip: Whether to disable GeoIP for this event.
+
+        Examples:
+            ```python
+            # $ai_generation event with blobs
+            posthog.capture_ai(
+                "$ai_generation",
+                distinct_id="user_123",
+                properties={
+                    "$ai_model": "gpt-4",
+                    "$ai_provider": "openai",
+                    "$ai_trace_id": "trace_abc123",
+                    "$ai_input": {
+                        "messages": [
+                            {"role": "user", "content": "Hello!"}
+                        ]
+                    },
+                    "$ai_output_choices": {
+                        "choices": [{"message": {"role": "assistant", "content": "Hi there!"}}]
+                    },
+                    "$ai_completion_tokens": 150,
+                    "$ai_prompt_tokens": 50
+                },
+                blob_properties=["$ai_input", "$ai_output_choices"]
+            )
+            ```
+
+            ```python
+            # $ai_trace event
+            posthog.capture_ai(
+                "$ai_trace",
+                distinct_id="user_123",
+                properties={
+                    "$ai_model": "gpt-4",
+                    "$ai_trace_id": "trace_abc123"
+                }
+            )
+            ```
+
+            ```python
+            # $ai_metric event
+            posthog.capture_ai(
+                "$ai_metric",
+                distinct_id="user_123",
+                properties={
+                    "$ai_model": "gpt-4",
+                    "$ai_trace_id": "trace_abc123",
+                    "$ai_metric_name": "accuracy",
+                    "$ai_metric_value": "0.95"
+                }
+            )
+            ```
+
+        Category:
+            AI Events
+
+        Note: This method sends events synchronously to the AI endpoint, bypassing the queue system.
+        """
+        import requests
+
+        if self.disabled:
+            return None
+
+        distinct_id = kwargs.get("distinct_id", None)
+        properties = kwargs.get("properties", None)
+        timestamp = kwargs.get("timestamp", None)
+        uuid = kwargs.get("uuid", None)
+        groups = kwargs.get("groups", None)
+        blob_properties = kwargs.get("blob_properties", None)
+
+        # Default blob properties based on event type if not provided
+        if blob_properties is None:
+            default_blob_properties = {
+                "$ai_generation": ["$ai_input", "$ai_output_choices"],
+                "$ai_trace": ["$ai_input_state", "$ai_output_state"],
+                "$ai_span": ["$ai_input_state", "$ai_output_state"],
+                "$ai_embedding": ["$ai_input"],
+            }
+            blob_properties = default_blob_properties.get(event, [])
+
+        properties = properties or {}
+
+        # Get distinct_id
+        (distinct_id, personless) = get_identity_state(distinct_id)
+
+        if personless and "$process_person_profile" not in properties:
+            properties["$process_person_profile"] = False
+
+        # Prepare timestamp
+        if timestamp is None:
+            timestamp = datetime.now(tz=tzutc())
+
+        timestamp = guess_timezone(timestamp)
+        timestamp_str = timestamp.isoformat()
+
+        # Generate UUID if not provided
+        if uuid:
+            uuid_str = stringify_id(uuid)
+        else:
+            uuid_str = stringify_id(uuid4())
+
+        # Add groups to properties
+        if groups:
+            properties["$groups"] = groups
+
+        # Prepare event data (the main event part)
+        event_data = {
+            "uuid": uuid_str,
+            "event": event,
+            "distinct_id": stringify_id(distinct_id),
+            "timestamp": timestamp_str,
+        }
+
+        # Build multipart body
+        body_bytes, content_type = _build_multipart_body(
+            event_data, properties, blob_properties
+        )
+
+        # Send request to AI endpoint
+        url = remove_trailing_slash(self.host) + "/i/v0/ai"
+
+        headers = {
+            "Content-Type": content_type,
+            "Authorization": f"Bearer {self.api_key}",
+            "User-Agent": f"posthog-python/{VERSION}",
+        }
+
+        self.log.debug(f"Sending AI event to {url}: {event} (uuid: {uuid_str})")
+
+        try:
+            response = requests.post(
+                url,
+                data=body_bytes,
+                headers=headers,
+                timeout=self.timeout,
+            )
+
+            if response.status_code == 200:
+                self.log.debug(f"AI event captured successfully: {event}")
+                return uuid_str
+            else:
+                error_message = f"AI capture failed with status {response.status_code}"
+                try:
+                    error_detail = response.json()
+                    error_message = f"{error_message}: {error_detail}"
+                except Exception:
+                    error_message = f"{error_message}: {response.text}"
+
+                self.log.error(error_message)
+
+                if self.debug:
+                    raise Exception(error_message)
+
+                return None
+
+        except Exception as e:
+            self.log.exception(f"Error sending AI event: {e}")
+
+            if self.debug:
+                raise e
+
+            return None
+
     def _enqueue(self, msg, disable_geoip):
         # type: (...) -> Optional[str]
         """Push a new `msg` onto the queue, return `(success, msg)`"""
@@ -1543,7 +1886,7 @@ class Client(object):
                     self.log.debug(
                         f"Successfully computed flag locally: {key} -> {response}"
                     )
-                except InconclusiveMatchError as e:
+                except (RequiresServerEvaluation, InconclusiveMatchError) as e:
                     self.log.debug(f"Failed to compute flag {key} locally: {e}")
                 except Exception as e:
                     self.log.exception(
@@ -1814,7 +2157,7 @@ class Client(object):
             )
         )
 
-        response, fallback_to_decide = self._get_all_flags_and_payloads_locally(
+        response, fallback_to_flags = self._get_all_flags_and_payloads_locally(
             distinct_id,
             groups=groups,
             person_properties=person_properties,
@@ -1822,7 +2165,7 @@ class Client(object):
             flag_keys_to_evaluate=flag_keys_to_evaluate,
         )
 
-        if fallback_to_decide and not only_evaluate_locally:
+        if fallback_to_flags and not only_evaluate_locally:
             try:
                 decide_response = self.get_flags_decision(
                     distinct_id,
@@ -1858,7 +2201,7 @@ class Client(object):
 
         flags: dict[str, FlagValue] = {}
         payloads: dict[str, str] = {}
-        fallback_to_decide = False
+        fallback_to_flags = False
         # If loading in previous line failed
         if self.feature_flags:
             # Filter flags based on flag_keys_to_evaluate if provided
@@ -1886,19 +2229,19 @@ class Client(object):
                         payloads[flag["key"]] = matched_payload
                 except InconclusiveMatchError:
                     # No need to log this, since it's just telling us to fall back to `/flags`
-                    fallback_to_decide = True
+                    fallback_to_flags = True
                 except Exception as e:
                     self.log.exception(
                         f"[FEATURE FLAGS] Error while computing variant and payload: {e}"
                     )
-                    fallback_to_decide = True
+                    fallback_to_flags = True
         else:
-            fallback_to_decide = True
+            fallback_to_flags = True
 
         return {
             "featureFlags": flags,
             "featureFlagPayloads": payloads,
-        }, fallback_to_decide
+        }, fallback_to_flags
 
     def _initialize_flag_cache(self, cache_url):
         """Initialize feature flag cache for graceful degradation during service outages.

posthog/feature_flags.py

@@ -22,6 +22,18 @@ class InconclusiveMatchError(Exception):
     pass
 
 
+class RequiresServerEvaluation(Exception):
+    """
+    Raised when feature flag evaluation requires server-side data that is not
+    available locally (e.g., static cohorts, experience continuity).
+
+    This error should propagate immediately to trigger API fallback, unlike
+    InconclusiveMatchError which allows trying other conditions.
+    """
+
+    pass
+
+
 # This function takes a distinct_id and a feature flag key and returns a float between 0 and 1.
 # Given the same distinct_id and key, it'll always return the same float. These floats are
 # uniformly distributed between 0 and 1, so if we want to show this feature to 20% of traffic
@@ -220,14 +232,7 @@ def match_feature_flag_properties(
     ) or []
     valid_variant_keys = [variant["key"] for variant in flag_variants]
 
-    # Stable sort conditions with variant overrides to the top. This ensures that if overrides are present, they are
-    # evaluated first, and the variant override is applied to the first matching condition.
-    sorted_flag_conditions = sorted(
-        flag_conditions,
-        key=lambda condition: 0 if condition.get("variant") else 1,
-    )
-
-    for condition in sorted_flag_conditions:
+    for condition in flag_conditions:
         try:
             # if any one condition resolves to True, we can shortcircuit and return
             # the matching variant
@@ -246,7 +251,12 @@ def match_feature_flag_properties(
                 else:
                     variant = get_matching_variant(flag, distinct_id)
                 return variant or True
+        except RequiresServerEvaluation:
+            # Static cohort or other missing server-side data - must fallback to API
+            raise
         except InconclusiveMatchError:
+            # Evaluation error (bad regex, invalid date, missing property, etc.)
+            # Track that we had an inconclusive match, but try other conditions
             is_inconclusive = True
 
     if is_inconclusive:
@@ -456,8 +466,8 @@ def match_cohort(
     # }
     cohort_id = str(property.get("value"))
     if cohort_id not in cohort_properties:
-        raise InconclusiveMatchError(
-            "can't match cohort without a given cohort property value"
+        raise RequiresServerEvaluation(
+            f"cohort {cohort_id} not found in local cohorts - likely a static cohort that requires server evaluation"
         )
 
     property_group = cohort_properties[cohort_id]
@@ -510,6 +520,9 @@ def match_property_group(
                     # OR group
                     if matches:
                         return True
+            except RequiresServerEvaluation:
+                # Immediately propagate - this condition requires server-side data
+                raise
             except InconclusiveMatchError as e:
                 log.debug(f"Failed to compute property {prop} locally: {e}")
                 error_matching_locally = True
@@ -559,6 +572,9 @@ def match_property_group(
                         return True
                     if not matches and negation:
                         return True
+            except RequiresServerEvaluation:
+                # Immediately propagate - this condition requires server-side data
+                raise
             except InconclusiveMatchError as e:
                 log.debug(f"Failed to compute property {prop} locally: {e}")
                 error_matching_locally = True

posthog/integrations/django.py

@@ -1,10 +1,24 @@
 from typing import TYPE_CHECKING, cast
-from posthog import contexts, capture_exception
+from posthog import contexts
 from posthog.client import Client
 
+try:
+    from asgiref.sync import iscoroutinefunction, markcoroutinefunction
+except ImportError:
+    # Fallback for older Django versions without asgiref
+    import asyncio
+
+    iscoroutinefunction = asyncio.iscoroutinefunction
+
+    # No-op fallback for markcoroutinefunction
+    # Older Django versions without asgiref typically don't support async middleware anyway
+    def markcoroutinefunction(func):
+        return func
+
+
 if TYPE_CHECKING:
     from django.http import HttpRequest, HttpResponse  # noqa: F401
-    from typing import Callable, Dict, Any, Optional  # noqa: F401
+    from typing import Callable, Dict, Any, Optional, Union, Awaitable  # noqa: F401
 
 
 class PosthogContextMiddleware:
@@ -31,11 +45,24 @@ class PosthogContextMiddleware:
     See the context documentation for more information. The extracted distinct ID and session ID, if found, are used to
     associate all events captured in the middleware context with the same distinct ID and session as currently active on the
     frontend. See the documentation for `set_context_session` and `identify_context` for more details.
+
+    This middleware is hybrid-capable: it supports both WSGI (sync) and ASGI (async) Django applications. The middleware
+    detects at initialization whether the next middleware in the chain is async or sync, and adapts its behavior accordingly.
+    This ensures compatibility with both pure sync and pure async middleware chains, as well as mixed chains in ASGI mode.
     """
 
+    sync_capable = True
+    async_capable = True
+
     def __init__(self, get_response):
-        # type: (Callable[[HttpRequest], HttpResponse]) -> None
+        # type: (Union[Callable[[HttpRequest], HttpResponse], Callable[[HttpRequest], Awaitable[HttpResponse]]]) -> None
         self.get_response = get_response
+        self._is_coroutine = iscoroutinefunction(get_response)
+
+        # Mark this instance as a coroutine function if get_response is async
+        # This is required for Django to correctly detect async middleware
+        if self._is_coroutine:
+            markcoroutinefunction(self)
 
         from django.conf import settings
 
@@ -158,24 +185,67 @@ class PosthogContextMiddleware:
         return user_id, email
 
     def __call__(self, request):
-        # type: (HttpRequest) -> HttpResponse
+        # type: (HttpRequest) -> Union[HttpResponse, Awaitable[HttpResponse]]
+        """
+        Unified entry point for both sync and async request handling.
+
+        When sync_capable and async_capable are both True, Django passes requests
+        without conversion. This method detects the mode and routes accordingly.
+        """
+        if self._is_coroutine:
+            return self.__acall__(request)
+        else:
+            # Synchronous path
+            if self.request_filter and not self.request_filter(request):
+                return self.get_response(request)
+
+            with contexts.new_context(self.capture_exceptions, client=self.client):
+                for k, v in self.extract_tags(request).items():
+                    contexts.tag(k, v)
+
+                return self.get_response(request)
+
+    async def __acall__(self, request):
+        # type: (HttpRequest) -> Awaitable[HttpResponse]
+        """
+        Asynchronous entry point for async request handling.
+
+        This method is called when the middleware chain is async.
+        """
         if self.request_filter and not self.request_filter(request):
-            return self.get_response(request)
+            return await self.get_response(request)
 
         with contexts.new_context(self.capture_exceptions, client=self.client):
             for k, v in self.extract_tags(request).items():
                 contexts.tag(k, v)
 
-            return self.get_response(request)
+            return await self.get_response(request)
 
     def process_exception(self, request, exception):
+        # type: (HttpRequest, Exception) -> None
+        """
+        Process exceptions from views and downstream middleware.
+
+        Django calls this WHILE still inside the context created by __call__,
+        so request tags have already been extracted and set. This method just
+        needs to capture the exception directly.
+
+        Django converts view exceptions into responses before they propagate through
+        the middleware stack, so the context manager in __call__/__acall__ never sees them.
+
+        Note: Django's process_exception is always synchronous, even for async views.
+        """
         if self.request_filter and not self.request_filter(request):
             return
 
         if not self.capture_exceptions:
             return
 
+        # Context and tags already set by __call__ or __acall__
+        # Just capture the exception
         if self.client:
             self.client.capture_exception(exception)
         else:
+            from posthog import capture_exception
+
             capture_exception(exception)

posthog/test/ai/anthropic/test_anthropic.py

@@ -1,5 +1,4 @@
 import os
-import time
 from unittest.mock import patch
 
 import pytest
@@ -13,14 +12,95 @@ try:
 except ImportError:
     ANTHROPIC_AVAILABLE = False
 
-ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY")
-
 # Skip all tests if Anthropic is not available
 pytestmark = pytest.mark.skipif(
     not ANTHROPIC_AVAILABLE, reason="Anthropic package is not available"
 )
 
 
+# =======================
+# Reusable Mock Helpers
+# =======================
+
+
+class MockContent:
+    """Reusable mock content class for Anthropic responses."""
+
+    def __init__(self, text="Bar", content_type="text"):
+        self.type = content_type
+        self.text = text
+
+
+class MockUsage:
+    """Reusable mock usage class for Anthropic responses."""
+
+    def __init__(
+        self,
+        input_tokens=18,
+        output_tokens=1,
+        cache_read_input_tokens=0,
+        cache_creation_input_tokens=0,
+    ):
+        self.input_tokens = input_tokens
+        self.output_tokens = output_tokens
+        self.cache_read_input_tokens = cache_read_input_tokens
+        self.cache_creation_input_tokens = cache_creation_input_tokens
+
+
+class MockResponse:
+    """Reusable mock response class for Anthropic messages."""
+
+    def __init__(
+        self,
+        content_text="Bar",
+        model="claude-3-opus-20240229",
+        input_tokens=18,
+        output_tokens=1,
+        cache_read=0,
+        cache_creation=0,
+    ):
+        self.content = [MockContent(text=content_text)]
+        self.model = model
+        self.usage = MockUsage(
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            cache_read_input_tokens=cache_read,
+            cache_creation_input_tokens=cache_creation,
+        )
+
+
+def create_mock_response(**kwargs):
+    """Factory function to create mock responses with custom parameters."""
+    return MockResponse(**kwargs)
+
+
+# Streaming mock helpers
+class MockStreamEvent:
+    """Reusable mock event class for streaming responses."""
+
+    def __init__(self, event_type=None, **kwargs):
+        self.type = event_type
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+
+
+class MockContentBlock:
+    """Reusable mock content block for streaming."""
+
+    def __init__(self, block_type, **kwargs):
+        self.type = block_type
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+
+
+class MockDelta:
+    """Reusable mock delta for streaming events."""
+
+    def __init__(self, **kwargs):
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+
+
 @pytest.fixture
 def mock_client():
     with patch("posthog.client.Client") as mock_client:
@@ -46,22 +126,77 @@ def mock_anthropic_response():
 
 
 @pytest.fixture
-def mock_anthropic_stream():
-    class MockStreamEvent:
-        def __init__(self, content, usage=None):
-            self.content = content
-            self.usage = usage
+def mock_anthropic_stream_with_tools():
+    """Mock stream events for tool calls."""
+
+    class MockMessage:
+        def __init__(self):
+            self.usage = MockUsage(
+                input_tokens=50,
+                cache_creation_input_tokens=0,
+                cache_read_input_tokens=5,
+            )
 
     def stream_generator():
-        yield MockStreamEvent("A")
-        yield MockStreamEvent("B")
-        yield MockStreamEvent(
-            "C",
-            usage=Usage(
-                input_tokens=20,
-                output_tokens=10,
-            ),
+        # Message start with usage
+        event = MockStreamEvent("message_start")
+        event.message = MockMessage()
+        yield event
+
+        # Text block start
+        event = MockStreamEvent("content_block_start")
+        event.content_block = MockContentBlock("text")
+        event.index = 0
+        yield event
+
+        # Text delta
+        event = MockStreamEvent("content_block_delta")
+        event.delta = MockDelta(text="I'll check the weather for you.")
+        event.index = 0
+        yield event
+
+        # Text block stop
+        event = MockStreamEvent("content_block_stop")
+        event.index = 0
+        yield event
+
+        # Tool use block start
+        event = MockStreamEvent("content_block_start")
+        event.content_block = MockContentBlock(
+            "tool_use", id="toolu_stream123", name="get_weather"
         )
+        event.index = 1
+        yield event
+
+        # Tool input delta 1
+        event = MockStreamEvent("content_block_delta")
+        event.delta = MockDelta(
+            type="input_json_delta", partial_json='{"location": "San'
+        )
+        event.index = 1
+        yield event
+
+        # Tool input delta 2
+        event = MockStreamEvent("content_block_delta")
+        event.delta = MockDelta(
+            type="input_json_delta", partial_json=' Francisco", "unit": "celsius"}'
+        )
+        event.index = 1
+        yield event
+
+        # Tool block stop
+        event = MockStreamEvent("content_block_stop")
+        event.index = 1
+        yield event
+
+        # Message delta with final usage
+        event = MockStreamEvent("message_delta")
+        event.usage = MockUsage(output_tokens=25)
+        yield event
+
+        # Message stop
+        event = MockStreamEvent("message_stop")
+        yield event
 
     return stream_generator()
 
@@ -174,83 +309,6 @@ def test_basic_completion(mock_client, mock_anthropic_response):
         assert isinstance(props["$ai_latency"], float)
 
 
-def test_streaming(mock_client, mock_anthropic_stream):
-    with patch(
-        "anthropic.resources.Messages.create", return_value=mock_anthropic_stream
-    ):
-        client = Anthropic(api_key="test-key", posthog_client=mock_client)
-        response = client.messages.create(
-            model="claude-3-opus-20240229",
-            messages=[{"role": "user", "content": "Hello"}],
-            stream=True,
-            posthog_distinct_id="test-id",
-            posthog_properties={"foo": "bar"},
-        )
-
-        # Consume the stream
-        chunks = list(response)
-        assert len(chunks) == 3
-        assert chunks[0].content == "A"
-        assert chunks[1].content == "B"
-        assert chunks[2].content == "C"
-
-        # Wait a bit to ensure the capture is called
-        time.sleep(0.1)
-        assert mock_client.capture.call_count == 1
-
-        call_args = mock_client.capture.call_args[1]
-        props = call_args["properties"]
-
-        assert call_args["distinct_id"] == "test-id"
-        assert call_args["event"] == "$ai_generation"
-        assert props["$ai_provider"] == "anthropic"
-        assert props["$ai_model"] == "claude-3-opus-20240229"
-        assert props["$ai_input"] == [{"role": "user", "content": "Hello"}]
-        assert props["$ai_output_choices"] == [{"role": "assistant", "content": "ABC"}]
-        assert props["$ai_input_tokens"] == 20
-        assert props["$ai_output_tokens"] == 10
-        assert isinstance(props["$ai_latency"], float)
-        assert props["foo"] == "bar"
-
-
-def test_streaming_with_stream_endpoint(mock_client, mock_anthropic_stream):
-    with patch(
-        "anthropic.resources.Messages.create", return_value=mock_anthropic_stream
-    ):
-        client = Anthropic(api_key="test-key", posthog_client=mock_client)
-        response = client.messages.stream(
-            model="claude-3-opus-20240229",
-            messages=[{"role": "user", "content": "Hello"}],
-            posthog_distinct_id="test-id",
-            posthog_properties={"foo": "bar"},
-        )
-
-        # Consume the stream
-        chunks = list(response)
-        assert len(chunks) == 3
-        assert chunks[0].content == "A"
-        assert chunks[1].content == "B"
-        assert chunks[2].content == "C"
-
-        # Wait a bit to ensure the capture is called
-        time.sleep(0.1)
-        assert mock_client.capture.call_count == 1
-
-        call_args = mock_client.capture.call_args[1]
-        props = call_args["properties"]
-
-        assert call_args["distinct_id"] == "test-id"
-        assert call_args["event"] == "$ai_generation"
-        assert props["$ai_provider"] == "anthropic"
-        assert props["$ai_model"] == "claude-3-opus-20240229"
-        assert props["$ai_input"] == [{"role": "user", "content": "Hello"}]
-        assert props["$ai_output_choices"] == [{"role": "assistant", "content": "ABC"}]
-        assert props["$ai_input_tokens"] == 20
-        assert props["$ai_output_tokens"] == 10
-        assert isinstance(props["$ai_latency"], float)
-        assert props["foo"] == "bar"
-
-
 def test_groups(mock_client, mock_anthropic_response):
     with patch(
         "anthropic.resources.Messages.create", return_value=mock_anthropic_response
@@ -313,18 +371,23 @@ def test_privacy_mode_global(mock_client, mock_anthropic_response):
         assert props["$ai_output_choices"] is None
 
 
-@pytest.mark.skipif(not ANTHROPIC_API_KEY, reason="ANTHROPIC_API_KEY is not set")
 def test_basic_integration(mock_client):
-    client = Anthropic(posthog_client=mock_client)
-    client.messages.create(
-        model="claude-3-opus-20240229",
-        messages=[{"role": "user", "content": "Foo"}],
-        max_tokens=1,
-        temperature=0,
-        posthog_distinct_id="test-id",
-        posthog_properties={"foo": "bar"},
-        system="You must always answer with 'Bar'.",
-    )
+    """Test basic non-streaming integration."""
+
+    with patch(
+        "anthropic.resources.Messages.create",
+        return_value=create_mock_response(),
+    ):
+        client = Anthropic(posthog_client=mock_client)
+        client.messages.create(
+            model="claude-3-opus-20240229",
+            messages=[{"role": "user", "content": "Foo"}],
+            max_tokens=1,
+            temperature=0,
+            posthog_distinct_id="test-id",
+            posthog_properties={"foo": "bar"},
+            system="You must always answer with 'Bar'.",
+        )
 
     assert mock_client.capture.call_count == 1
 
@@ -349,17 +412,28 @@ def test_basic_integration(mock_client):
     assert isinstance(props["$ai_latency"], float)
 
 
-@pytest.mark.skipif(not ANTHROPIC_API_KEY, reason="ANTHROPIC_API_KEY is not set")
 async def test_basic_async_integration(mock_client):
-    client = AsyncAnthropic(posthog_client=mock_client)
-    await client.messages.create(
-        model="claude-3-opus-20240229",
-        messages=[{"role": "user", "content": "You must always answer with 'Bar'."}],
-        max_tokens=1,
-        temperature=0,
-        posthog_distinct_id="test-id",
-        posthog_properties={"foo": "bar"},
-    )
+    """Test async non-streaming integration."""
+
+    # Make the mock async
+    async def mock_async_create(**kwargs):
+        return create_mock_response(input_tokens=16)
+
+    with patch(
+        "anthropic.resources.messages.AsyncMessages.create",
+        side_effect=mock_async_create,
+    ):
+        client = AsyncAnthropic(posthog_client=mock_client)
+        await client.messages.create(
+            model="claude-3-opus-20240229",
+            messages=[
+                {"role": "user", "content": "You must always answer with 'Bar'."}
+            ],
+            max_tokens=1,
+            temperature=0,
+            posthog_distinct_id="test-id",
+            posthog_properties={"foo": "bar"},
+        )
 
     assert mock_client.capture.call_count == 1
 
@@ -381,52 +455,50 @@ async def test_basic_async_integration(mock_client):
     assert isinstance(props["$ai_latency"], float)
 
 
-def test_streaming_system_prompt(mock_client, mock_anthropic_stream):
+async def test_async_streaming_system_prompt(mock_client):
+    """Test async streaming with system prompt."""
+
+    # Create a simple mock async stream using reusable helpers
+    async def mock_async_stream():
+        # Yield some events
+        yield MockStreamEvent(type="message_start")
+        yield MockStreamEvent(type="content_block_start")
+        yield MockStreamEvent(type="content_block_delta", text="Bar")
+
+        # Final message with usage
+        final_msg = MockStreamEvent(type="message_delta")
+        final_msg.usage = MockUsage(
+            input_tokens=10,
+            output_tokens=5,
+            cache_read_input_tokens=0,
+            cache_creation_input_tokens=0,
+        )
+        yield final_msg
+
+    # Mock create to return a coroutine that yields the async generator
+    # This matches the actual behavior when stream=True with await
+    async def async_create_wrapper(**kwargs):
+        return mock_async_stream()
+
     with patch(
-        "anthropic.resources.Messages.create", return_value=mock_anthropic_stream
+        "anthropic.resources.messages.AsyncMessages.create",
+        side_effect=async_create_wrapper,
     ):
-        client = Anthropic(api_key="test-key", posthog_client=mock_client)
-        response = client.messages.create(
+        client = AsyncAnthropic(posthog_client=mock_client)
+        response = await client.messages.create(
             model="claude-3-opus-20240229",
-            system="Foo",
-            messages=[{"role": "user", "content": "Bar"}],
+            system="You must always answer with 'Bar'.",
+            messages=[{"role": "user", "content": "Foo"}],
             stream=True,
+            max_tokens=1,
         )
 
-        # Consume the stream
-        list(response)
+        # Consume the stream - async finally block completes before this returns
+        [c async for c in response]
 
-        # Wait a bit to ensure the capture is called
-        time.sleep(0.1)
+        # Capture happens in the async finally block before generator completes
         assert mock_client.capture.call_count == 1
 
-        call_args = mock_client.capture.call_args[1]
-        props = call_args["properties"]
-
-        assert props["$ai_input"] == [
-            {"role": "system", "content": "Foo"},
-            {"role": "user", "content": "Bar"},
-        ]
-
-
-@pytest.mark.skipif(not ANTHROPIC_API_KEY, reason="ANTHROPIC_API_KEY is not set")
-async def test_async_streaming_system_prompt(mock_client, mock_anthropic_stream):
-    client = AsyncAnthropic(posthog_client=mock_client)
-    response = await client.messages.create(
-        model="claude-3-opus-20240229",
-        system="You must always answer with 'Bar'.",
-        messages=[{"role": "user", "content": "Foo"}],
-        stream=True,
-        max_tokens=1,
-    )
-
-    # Consume the stream
-    [c async for c in response]
-
-    # Wait a bit to ensure the capture is called
-    time.sleep(0.1)
-    assert mock_client.capture.call_count == 1
-
     call_args = mock_client.capture.call_args[1]
     props = call_args["properties"]
 
@@ -746,3 +818,219 @@ def test_async_tool_calls_in_output_choices(
         assert props["$ai_input_tokens"] == 25
         assert props["$ai_output_tokens"] == 15
         assert props["$ai_http_status"] == 200
+
+
+def test_streaming_with_tool_calls(mock_client, mock_anthropic_stream_with_tools):
+    """Test that tool calls are properly captured in streaming mode."""
+    with patch(
+        "anthropic.resources.Messages.create",
+        return_value=mock_anthropic_stream_with_tools,
+    ):
+        client = Anthropic(api_key="test-key", posthog_client=mock_client)
+        response = client.messages.create(
+            model="claude-3-5-sonnet-20241022",
+            system="You are a helpful weather assistant.",
+            messages=[
+                {"role": "user", "content": "What's the weather in San Francisco?"}
+            ],
+            tools=[
+                {
+                    "name": "get_weather",
+                    "description": "Get weather information",
+                    "input_schema": {
+                        "type": "object",
+                        "properties": {
+                            "location": {"type": "string"},
+                            "unit": {"type": "string"},
+                        },
+                        "required": ["location"],
+                    },
+                }
+            ],
+            stream=True,
+            posthog_distinct_id="test-id",
+        )
+
+        # Consume the stream - this triggers the finally block synchronously
+        list(response)
+
+        # Capture happens synchronously when generator is exhausted
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+
+        assert call_args["distinct_id"] == "test-id"
+        assert call_args["event"] == "$ai_generation"
+        assert props["$ai_provider"] == "anthropic"
+        assert props["$ai_model"] == "claude-3-5-sonnet-20241022"
+
+        # Verify system prompt is included in input
+        assert props["$ai_input"] == [
+            {"role": "system", "content": "You are a helpful weather assistant."},
+            {"role": "user", "content": "What's the weather in San Francisco?"},
+        ]
+
+        # Verify that tools are captured in the properties
+        assert props["$ai_tools"] == [
+            {
+                "name": "get_weather",
+                "description": "Get weather information",
+                "input_schema": {
+                    "type": "object",
+                    "properties": {
+                        "location": {"type": "string"},
+                        "unit": {"type": "string"},
+                    },
+                    "required": ["location"],
+                },
+            }
+        ]
+
+        # Verify output contains both text and tool call
+        output_choices = props["$ai_output_choices"]
+        assert len(output_choices) == 1
+
+        assistant_message = output_choices[0]
+        assert assistant_message["role"] == "assistant"
+
+        content = assistant_message["content"]
+        assert isinstance(content, list)
+        assert len(content) == 2
+
+        # Verify text block
+        text_block = content[0]
+        assert text_block["type"] == "text"
+        assert text_block["text"] == "I'll check the weather for you."
+
+        # Verify tool call block
+        tool_block = content[1]
+        assert tool_block["type"] == "function"
+        assert tool_block["id"] == "toolu_stream123"
+        assert tool_block["function"]["name"] == "get_weather"
+        assert tool_block["function"]["arguments"] == {
+            "location": "San Francisco",
+            "unit": "celsius",
+        }
+
+        # Check token usage
+        assert props["$ai_input_tokens"] == 50
+        assert props["$ai_output_tokens"] == 25
+        assert props["$ai_cache_read_input_tokens"] == 5
+        assert props["$ai_cache_creation_input_tokens"] == 0
+
+
+def test_async_streaming_with_tool_calls(mock_client, mock_anthropic_stream_with_tools):
+    """Test that tool calls are properly captured in async streaming mode."""
+    import asyncio
+
+    async def mock_async_generator():
+        # Convert regular generator to async generator
+        for event in mock_anthropic_stream_with_tools:
+            yield event
+
+    async def mock_async_create(**kwargs):
+        # Return the async generator (to be awaited by the implementation)
+        return mock_async_generator()
+
+    with patch(
+        "anthropic.resources.AsyncMessages.create",
+        side_effect=mock_async_create,
+    ):
+        async_client = AsyncAnthropic(api_key="test-key", posthog_client=mock_client)
+
+        async def run_test():
+            response = await async_client.messages.create(
+                model="claude-3-5-sonnet-20241022",
+                system="You are a helpful weather assistant.",
+                messages=[
+                    {"role": "user", "content": "What's the weather in San Francisco?"}
+                ],
+                tools=[
+                    {
+                        "name": "get_weather",
+                        "description": "Get weather information",
+                        "input_schema": {
+                            "type": "object",
+                            "properties": {
+                                "location": {"type": "string"},
+                                "unit": {"type": "string"},
+                            },
+                            "required": ["location"],
+                        },
+                    }
+                ],
+                stream=True,
+                posthog_distinct_id="test-id",
+            )
+
+            # Consume the async stream
+            [event async for event in response]
+
+        # asyncio.run() waits for all async operations to complete
+        asyncio.run(run_test())
+
+        # Capture completes before asyncio.run() returns
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+
+        assert call_args["distinct_id"] == "test-id"
+        assert call_args["event"] == "$ai_generation"
+        assert props["$ai_provider"] == "anthropic"
+        assert props["$ai_model"] == "claude-3-5-sonnet-20241022"
+
+        # Verify system prompt is included in input
+        assert props["$ai_input"] == [
+            {"role": "system", "content": "You are a helpful weather assistant."},
+            {"role": "user", "content": "What's the weather in San Francisco?"},
+        ]
+
+        # Verify that tools are captured in the properties
+        assert props["$ai_tools"] == [
+            {
+                "name": "get_weather",
+                "description": "Get weather information",
+                "input_schema": {
+                    "type": "object",
+                    "properties": {
+                        "location": {"type": "string"},
+                        "unit": {"type": "string"},
+                    },
+                    "required": ["location"],
+                },
+            }
+        ]
+
+        # Verify output contains both text and tool call
+        output_choices = props["$ai_output_choices"]
+        assert len(output_choices) == 1
+
+        assistant_message = output_choices[0]
+        assert assistant_message["role"] == "assistant"
+
+        content = assistant_message["content"]
+        assert isinstance(content, list)
+        assert len(content) == 2
+
+        # Verify text block
+        text_block = content[0]
+        assert text_block["type"] == "text"
+        assert text_block["text"] == "I'll check the weather for you."
+
+        # Verify tool call block
+        tool_block = content[1]
+        assert tool_block["type"] == "function"
+        assert tool_block["id"] == "toolu_stream123"
+        assert tool_block["function"]["name"] == "get_weather"
+        assert tool_block["function"]["arguments"] == {
+            "location": "San Francisco",
+            "unit": "celsius",
+        }
+
+        # Check token usage
+        assert props["$ai_input_tokens"] == 50
+        assert props["$ai_output_tokens"] == 25
+        assert props["$ai_cache_read_input_tokens"] == 5
+        assert props["$ai_cache_creation_input_tokens"] == 0

posthog/test/ai/gemini/test_gemini.py

@@ -31,6 +31,9 @@ def mock_gemini_response():
     mock_usage = MagicMock()
     mock_usage.prompt_token_count = 20
     mock_usage.candidates_token_count = 10
+    # Ensure cache and reasoning tokens are not present (not MagicMock)
+    mock_usage.cached_content_token_count = 0
+    mock_usage.thoughts_token_count = 0
     mock_response.usage_metadata = mock_usage
 
     mock_candidate = MagicMock()
@@ -64,6 +67,8 @@ def mock_gemini_response_with_function_calls():
     mock_usage = MagicMock()
     mock_usage.prompt_token_count = 25
     mock_usage.candidates_token_count = 15
+    mock_usage.cached_content_token_count = 0
+    mock_usage.thoughts_token_count = 0
     mock_response.usage_metadata = mock_usage
 
     # Mock function call
@@ -110,6 +115,8 @@ def mock_gemini_response_function_calls_only():
     mock_usage = MagicMock()
     mock_usage.prompt_token_count = 30
     mock_usage.candidates_token_count = 12
+    mock_usage.cached_content_token_count = 0
+    mock_usage.thoughts_token_count = 0
     mock_response.usage_metadata = mock_usage
 
     # Mock function call
@@ -180,6 +187,8 @@ def test_new_client_streaming_with_generate_content_stream(
         mock_usage1 = MagicMock()
         mock_usage1.prompt_token_count = 10
         mock_usage1.candidates_token_count = 5
+        mock_usage1.cached_content_token_count = 0
+        mock_usage1.thoughts_token_count = 0
         mock_chunk1.usage_metadata = mock_usage1
 
         mock_chunk2 = MagicMock()
@@ -187,6 +196,8 @@ def test_new_client_streaming_with_generate_content_stream(
         mock_usage2 = MagicMock()
         mock_usage2.prompt_token_count = 10
         mock_usage2.candidates_token_count = 10
+        mock_usage2.cached_content_token_count = 0
+        mock_usage2.thoughts_token_count = 0
         mock_chunk2.usage_metadata = mock_usage2
 
         yield mock_chunk1
@@ -226,6 +237,91 @@ def test_new_client_streaming_with_generate_content_stream(
     assert isinstance(props["$ai_latency"], float)
 
 
+def test_new_client_streaming_with_tools(mock_client, mock_google_genai_client):
+    """Test that tools are captured in streaming mode"""
+
+    def mock_streaming_response():
+        mock_chunk1 = MagicMock()
+        mock_chunk1.text = "I'll check "
+        mock_usage1 = MagicMock()
+        mock_usage1.prompt_token_count = 15
+        mock_usage1.candidates_token_count = 5
+        mock_usage1.cached_content_token_count = 0
+        mock_usage1.thoughts_token_count = 0
+        mock_chunk1.usage_metadata = mock_usage1
+
+        mock_chunk2 = MagicMock()
+        mock_chunk2.text = "the weather"
+        mock_usage2 = MagicMock()
+        mock_usage2.prompt_token_count = 15
+        mock_usage2.candidates_token_count = 10
+        mock_usage2.cached_content_token_count = 0
+        mock_usage2.thoughts_token_count = 0
+        mock_chunk2.usage_metadata = mock_usage2
+
+        yield mock_chunk1
+        yield mock_chunk2
+
+    # Mock the generate_content_stream method
+    mock_google_genai_client.models.generate_content_stream.return_value = (
+        mock_streaming_response()
+    )
+
+    client = Client(api_key="test-key", posthog_client=mock_client)
+
+    # Create mock tools configuration
+    mock_tool = MagicMock()
+    mock_tool.function_declarations = [
+        MagicMock(
+            name="get_current_weather",
+            description="Gets the current weather for a given location.",
+            parameters=MagicMock(
+                type="OBJECT",
+                properties={
+                    "location": MagicMock(
+                        type="STRING",
+                        description="The city and state, e.g. San Francisco, CA",
+                    )
+                },
+                required=["location"],
+            ),
+        )
+    ]
+
+    mock_config = MagicMock()
+    mock_config.tools = [mock_tool]
+
+    response = client.models.generate_content_stream(
+        model="gemini-2.0-flash",
+        contents=["What's the weather in SF?"],
+        config=mock_config,
+        posthog_distinct_id="test-id",
+        posthog_properties={"feature": "streaming_with_tools"},
+    )
+
+    chunks = list(response)
+    assert len(chunks) == 2
+    assert chunks[0].text == "I'll check "
+    assert chunks[1].text == "the weather"
+
+    # Check that the streaming event was captured with tools
+    assert mock_client.capture.call_count == 1
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+
+    assert call_args["distinct_id"] == "test-id"
+    assert call_args["event"] == "$ai_generation"
+    assert props["$ai_provider"] == "gemini"
+    assert props["$ai_model"] == "gemini-2.0-flash"
+    assert props["$ai_input_tokens"] == 15
+    assert props["$ai_output_tokens"] == 10
+    assert props["feature"] == "streaming_with_tools"
+    assert isinstance(props["$ai_latency"], float)
+
+    # Verify that tools are captured in the $ai_tools property in streaming mode
+    assert props["$ai_tools"] == [mock_tool]
+
+
 def test_new_client_groups(mock_client, mock_google_genai_client, mock_gemini_response):
     """Test groups functionality with new Client API"""
     mock_google_genai_client.models.generate_content.return_value = mock_gemini_response
@@ -302,12 +398,32 @@ def test_new_client_different_input_formats(
     props = call_args["properties"]
     assert props["$ai_input"] == [{"role": "user", "content": "Hello"}]
 
-    # Test list input
-    mock_client.capture.reset_mock()
-    mock_part = MagicMock()
-    mock_part.text = "List item"
+    # Test Gemini-specific format with parts array (like in the screenshot)
+    mock_client.reset_mock()
     client.models.generate_content(
-        model="gemini-2.0-flash", contents=[mock_part], posthog_distinct_id="test-id"
+        model="gemini-2.0-flash",
+        contents=[{"role": "user", "parts": [{"text": "hey"}]}],
+        posthog_distinct_id="test-id",
+    )
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+    assert props["$ai_input"] == [{"role": "user", "content": "hey"}]
+
+    # Test multiple parts in the parts array
+    mock_client.reset_mock()
+    client.models.generate_content(
+        model="gemini-2.0-flash",
+        contents=[{"role": "user", "parts": [{"text": "Hello "}, {"text": "world"}]}],
+        posthog_distinct_id="test-id",
+    )
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+    assert props["$ai_input"] == [{"role": "user", "content": "Hello world"}]
+
+    # Test list input with string
+    mock_client.capture.reset_mock()
+    client.models.generate_content(
+        model="gemini-2.0-flash", contents=["List item"], posthog_distinct_id="test-id"
     )
     call_args = mock_client.capture.call_args[1]
     props = call_args["properties"]
@@ -500,6 +616,8 @@ def test_tool_use_response(mock_client, mock_google_genai_client, mock_gemini_re
 
     mock_config = MagicMock()
     mock_config.tools = [mock_tool]
+    # Explicitly specify this config doesn't have system_instruction
+    del mock_config.system_instruction
 
     response = client.models.generate_content(
         model="gemini-2.5-flash",
@@ -629,3 +747,93 @@ def test_function_calls_only_no_content(
     assert props["$ai_input_tokens"] == 30
     assert props["$ai_output_tokens"] == 12
     assert props["$ai_http_status"] == 200
+
+
+def test_cache_and_reasoning_tokens(mock_client, mock_google_genai_client):
+    """Test that cache and reasoning tokens are properly extracted"""
+    # Create a mock response with cache and reasoning tokens
+    mock_response = MagicMock()
+    mock_response.text = "Test response with cache"
+
+    mock_usage = MagicMock()
+    mock_usage.prompt_token_count = 100
+    mock_usage.candidates_token_count = 50
+    mock_usage.cached_content_token_count = 30  # Cache tokens
+    mock_usage.thoughts_token_count = 10  # Reasoning tokens
+    mock_response.usage_metadata = mock_usage
+
+    # Mock candidates
+    mock_candidate = MagicMock()
+    mock_candidate.text = "Test response with cache"
+    mock_response.candidates = [mock_candidate]
+
+    mock_google_genai_client.models.generate_content.return_value = mock_response
+
+    client = Client(api_key="test-key", posthog_client=mock_client)
+
+    response = client.models.generate_content(
+        model="gemini-2.5-pro",
+        contents="Test with cache",
+        posthog_distinct_id="test-id",
+    )
+
+    assert response == mock_response
+    assert mock_client.capture.call_count == 1
+
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+
+    # Check that all token types are present
+    assert props["$ai_input_tokens"] == 100
+    assert props["$ai_output_tokens"] == 50
+    assert props["$ai_cache_read_input_tokens"] == 30
+    assert props["$ai_reasoning_tokens"] == 10
+
+
+def test_streaming_cache_and_reasoning_tokens(mock_client, mock_google_genai_client):
+    """Test that cache and reasoning tokens are properly extracted in streaming"""
+    # Create mock chunks with cache and reasoning tokens
+    chunk1 = MagicMock()
+    chunk1.text = "Hello "
+    chunk1_usage = MagicMock()
+    chunk1_usage.prompt_token_count = 100
+    chunk1_usage.candidates_token_count = 5
+    chunk1_usage.cached_content_token_count = 30  # Cache tokens
+    chunk1_usage.thoughts_token_count = 0
+    chunk1.usage_metadata = chunk1_usage
+
+    chunk2 = MagicMock()
+    chunk2.text = "world!"
+    chunk2_usage = MagicMock()
+    chunk2_usage.prompt_token_count = 100
+    chunk2_usage.candidates_token_count = 10
+    chunk2_usage.cached_content_token_count = 30  # Same cache tokens
+    chunk2_usage.thoughts_token_count = 5  # Reasoning tokens
+    chunk2.usage_metadata = chunk2_usage
+
+    mock_stream = iter([chunk1, chunk2])
+    mock_google_genai_client.models.generate_content_stream.return_value = mock_stream
+
+    client = Client(api_key="test-key", posthog_client=mock_client)
+
+    response = client.models.generate_content_stream(
+        model="gemini-2.5-pro",
+        contents="Test streaming with cache",
+        posthog_distinct_id="test-id",
+    )
+
+    # Consume the stream
+    result = list(response)
+    assert len(result) == 2
+
+    # Check PostHog capture was called
+    assert mock_client.capture.call_count == 1
+
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+
+    # Check that all token types are present (should use final chunk's usage)
+    assert props["$ai_input_tokens"] == 100
+    assert props["$ai_output_tokens"] == 10
+    assert props["$ai_cache_read_input_tokens"] == 30
+    assert props["$ai_reasoning_tokens"] == 5

posthog/test/ai/langchain/test_callbacks.py

@@ -204,6 +204,7 @@ def test_basic_chat_chain(mock_client, stream):
     # Generation is second
     assert generation_args["event"] == "$ai_generation"
     assert "distinct_id" in generation_args
+    assert generation_props["$ai_framework"] == "langchain"
     assert "$ai_model" in generation_props
     assert "$ai_provider" in generation_props
     assert generation_props["$ai_input"] == [

posthog/test/ai/openai/test_openai.py

@@ -1,5 +1,5 @@
 import time
-from unittest.mock import patch
+from unittest.mock import AsyncMock, patch
 
 import pytest
 
@@ -34,6 +34,7 @@ try:
     )
 
     from posthog.ai.openai import OpenAI
+    from posthog.ai.openai.openai_async import AsyncOpenAI
 
     OPENAI_AVAILABLE = True
 except ImportError:
@@ -218,6 +219,106 @@ def mock_openai_response_with_cached_tokens():
     )
 
 
+@pytest.fixture
+def streaming_tool_call_chunks():
+    return [
+        ChatCompletionChunk(
+            id="chunk1",
+            model="gpt-4",
+            object="chat.completion.chunk",
+            created=1234567890,
+            choices=[
+                ChoiceChunk(
+                    index=0,
+                    delta=ChoiceDelta(
+                        role="assistant",
+                        tool_calls=[
+                            ChoiceDeltaToolCall(
+                                index=0,
+                                id="call_abc123",
+                                type="function",
+                                function=ChoiceDeltaToolCallFunction(
+                                    name="get_weather",
+                                    arguments='{"location": "',
+                                ),
+                            )
+                        ],
+                    ),
+                    finish_reason=None,
+                )
+            ],
+        ),
+        ChatCompletionChunk(
+            id="chunk2",
+            model="gpt-4",
+            object="chat.completion.chunk",
+            created=1234567891,
+            choices=[
+                ChoiceChunk(
+                    index=0,
+                    delta=ChoiceDelta(
+                        tool_calls=[
+                            ChoiceDeltaToolCall(
+                                index=0,
+                                id="call_abc123",
+                                type="function",
+                                function=ChoiceDeltaToolCallFunction(
+                                    arguments='San Francisco"',
+                                ),
+                            )
+                        ],
+                    ),
+                    finish_reason=None,
+                )
+            ],
+        ),
+        ChatCompletionChunk(
+            id="chunk3",
+            model="gpt-4",
+            object="chat.completion.chunk",
+            created=1234567892,
+            choices=[
+                ChoiceChunk(
+                    index=0,
+                    delta=ChoiceDelta(
+                        tool_calls=[
+                            ChoiceDeltaToolCall(
+                                index=0,
+                                id="call_abc123",
+                                type="function",
+                                function=ChoiceDeltaToolCallFunction(
+                                    arguments=', "unit": "celsius"}',
+                                ),
+                            )
+                        ],
+                    ),
+                    finish_reason=None,
+                )
+            ],
+        ),
+        ChatCompletionChunk(
+            id="chunk4",
+            model="gpt-4",
+            object="chat.completion.chunk",
+            created=1234567893,
+            choices=[
+                ChoiceChunk(
+                    index=0,
+                    delta=ChoiceDelta(
+                        content="The weather in San Francisco is 15°C.",
+                    ),
+                    finish_reason=None,
+                )
+            ],
+            usage=CompletionUsage(
+                prompt_tokens=20,
+                completion_tokens=15,
+                total_tokens=35,
+            ),
+        ),
+    ]
+
+
 @pytest.fixture
 def mock_openai_response_with_tool_calls():
     return ChatCompletion(
@@ -734,109 +835,11 @@ def test_responses_api_tool_calls(mock_client, mock_responses_api_with_tool_call
         assert props["$ai_http_status"] == 200
 
 
-def test_streaming_with_tool_calls(mock_client):
-    # Create mock tool call chunks that will be returned in sequence
-    tool_call_chunks = [
-        ChatCompletionChunk(
-            id="chunk1",
-            model="gpt-4",
-            object="chat.completion.chunk",
-            created=1234567890,
-            choices=[
-                ChoiceChunk(
-                    index=0,
-                    delta=ChoiceDelta(
-                        role="assistant",
-                        tool_calls=[
-                            ChoiceDeltaToolCall(
-                                index=0,
-                                id="call_abc123",
-                                type="function",
-                                function=ChoiceDeltaToolCallFunction(
-                                    name="get_weather",
-                                    arguments='{"location": "',
-                                ),
-                            )
-                        ],
-                    ),
-                    finish_reason=None,
-                )
-            ],
-        ),
-        ChatCompletionChunk(
-            id="chunk2",
-            model="gpt-4",
-            object="chat.completion.chunk",
-            created=1234567891,
-            choices=[
-                ChoiceChunk(
-                    index=0,
-                    delta=ChoiceDelta(
-                        tool_calls=[
-                            ChoiceDeltaToolCall(
-                                index=0,
-                                id="call_abc123",
-                                type="function",
-                                function=ChoiceDeltaToolCallFunction(
-                                    arguments='San Francisco"',
-                                ),
-                            )
-                        ],
-                    ),
-                    finish_reason=None,
-                )
-            ],
-        ),
-        ChatCompletionChunk(
-            id="chunk3",
-            model="gpt-4",
-            object="chat.completion.chunk",
-            created=1234567892,
-            choices=[
-                ChoiceChunk(
-                    index=0,
-                    delta=ChoiceDelta(
-                        tool_calls=[
-                            ChoiceDeltaToolCall(
-                                index=0,
-                                id="call_abc123",
-                                type="function",
-                                function=ChoiceDeltaToolCallFunction(
-                                    arguments=', "unit": "celsius"}',
-                                ),
-                            )
-                        ],
-                    ),
-                    finish_reason=None,
-                )
-            ],
-        ),
-        ChatCompletionChunk(
-            id="chunk4",
-            model="gpt-4",
-            object="chat.completion.chunk",
-            created=1234567893,
-            choices=[
-                ChoiceChunk(
-                    index=0,
-                    delta=ChoiceDelta(
-                        content="The weather in San Francisco is 15°C.",
-                    ),
-                    finish_reason=None,
-                )
-            ],
-            usage=CompletionUsage(
-                prompt_tokens=20,
-                completion_tokens=15,
-                total_tokens=35,
-            ),
-        ),
-    ]
-
+def test_streaming_with_tool_calls(mock_client, streaming_tool_call_chunks):
     # Mock the create method to return our chunks
     with patch("openai.resources.chat.completions.Completions.create") as mock_create:
         # Set up the mock to return our chunks when iterated
-        mock_create.return_value = tool_call_chunks
+        mock_create.return_value = streaming_tool_call_chunks
 
         client = OpenAI(api_key="test-key", posthog_client=mock_client)
 
@@ -865,7 +868,7 @@ def test_streaming_with_tool_calls(mock_client):
 
         # Verify the chunks were returned correctly
         assert len(chunks) == 4
-        assert chunks == tool_call_chunks
+        assert chunks == streaming_tool_call_chunks
 
         # Verify the capture was called with the right arguments
         assert mock_client.capture.call_count == 1
@@ -890,10 +893,29 @@ def test_streaming_with_tool_calls(mock_client):
         assert defined_tool["function"]["description"] == "Get weather"
         assert defined_tool["function"]["parameters"] == {}
 
-        # Check that the content was also accumulated
+        # Check that both text content and tool calls were accumulated
+        output_content = props["$ai_output_choices"][0]["content"]
+
+        # Find text content and tool call in the output
+        text_content = None
+        tool_call_content = None
+        for item in output_content:
+            if item["type"] == "text":
+                text_content = item
+            elif item["type"] == "function":
+                tool_call_content = item
+
+        # Verify text content
+        assert text_content is not None
+        assert text_content["text"] == "The weather in San Francisco is 15°C."
+
+        # Verify tool call was captured
+        assert tool_call_content is not None
+        assert tool_call_content["id"] == "call_abc123"
+        assert tool_call_content["function"]["name"] == "get_weather"
         assert (
-            props["$ai_output_choices"][0]["content"]
-            == "The weather in San Francisco is 15°C."
+            tool_call_content["function"]["arguments"]
+            == '{"location": "San Francisco", "unit": "celsius"}'
         )
 
         # Check token usage
@@ -1014,6 +1036,248 @@ def test_responses_parse(mock_client, mock_parsed_response):
         assert isinstance(props["$ai_latency"], float)
 
 
+def test_responses_api_streaming_with_tokens(mock_client):
+    """Test that Responses API streaming properly captures token usage from response.usage."""
+    from openai.types.responses import ResponseUsage
+    from unittest.mock import MagicMock
+
+    # Create mock response chunks with usage data in the correct location
+    chunks = []
+
+    # First chunk - just content, no usage
+    chunk1 = MagicMock()
+    chunk1.type = "response.text.delta"
+    chunk1.text = "Test "
+    chunks.append(chunk1)
+
+    # Second chunk - more content
+    chunk2 = MagicMock()
+    chunk2.type = "response.text.delta"
+    chunk2.text = "response"
+    chunks.append(chunk2)
+
+    # Final chunk - completed event with usage in response.usage
+    chunk3 = MagicMock()
+    chunk3.type = "response.completed"
+    chunk3.response = MagicMock()
+    chunk3.response.usage = ResponseUsage(
+        input_tokens=25,
+        output_tokens=30,
+        total_tokens=55,
+        input_tokens_details={"prompt_tokens": 25, "cached_tokens": 0},
+        output_tokens_details={"reasoning_tokens": 0},
+    )
+    chunk3.response.output = ["Test response"]
+    chunks.append(chunk3)
+
+    captured_kwargs = {}
+
+    def mock_streaming_response(**kwargs):
+        # Capture the kwargs to verify stream_options was NOT added
+        captured_kwargs.update(kwargs)
+        return iter(chunks)
+
+    with patch(
+        "openai.resources.responses.Responses.create",
+        side_effect=mock_streaming_response,
+    ):
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+
+        # Consume the streaming response
+        response = client.responses.create(
+            model="gpt-4o-mini",
+            input=[{"role": "user", "content": "Test message"}],
+            stream=True,
+            posthog_distinct_id="test-id",
+            posthog_properties={"test": "streaming"},
+        )
+
+        # Consume all chunks
+        list(response)
+
+        # Verify stream_options was NOT added (Responses API doesn't support it)
+        assert "stream_options" not in captured_kwargs
+
+        # Verify capture was called
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+
+        # Verify tokens are captured correctly from response.usage (not 0)
+        assert call_args["distinct_id"] == "test-id"
+        assert call_args["event"] == "$ai_generation"
+        assert props["$ai_provider"] == "openai"
+        assert props["$ai_model"] == "gpt-4o-mini"
+        assert props["$ai_input_tokens"] == 25  # Should not be 0
+        assert props["$ai_output_tokens"] == 30  # Should not be 0
+        assert props["test"] == "streaming"
+        assert isinstance(props["$ai_latency"], float)
+
+
+@pytest.mark.asyncio
+async def test_async_chat_streaming_with_tool_calls(
+    mock_client, streaming_tool_call_chunks
+):
+    captured_kwargs = {}
+
+    async def mock_create(self, **kwargs):
+        captured_kwargs["kwargs"] = kwargs
+
+        async def chunk_iterable():
+            for chunk in streaming_tool_call_chunks:
+                yield chunk
+
+        return chunk_iterable()
+
+    with patch(
+        "openai.resources.chat.completions.AsyncCompletions.create", new=mock_create
+    ):
+        client = AsyncOpenAI(api_key="test-key", posthog_client=mock_client)
+
+        response_stream = await client.chat.completions.create(
+            model="gpt-4",
+            messages=[
+                {"role": "user", "content": "What's the weather in San Francisco?"}
+            ],
+            tools=[
+                {
+                    "type": "function",
+                    "function": {
+                        "name": "get_weather",
+                        "description": "Get weather",
+                        "parameters": {},
+                    },
+                }
+            ],
+            stream=True,
+            posthog_distinct_id="test-id",
+        )
+
+        chunks = []
+        async for chunk in response_stream:
+            chunks.append(chunk)
+
+    kwargs = captured_kwargs["kwargs"]
+    assert kwargs["stream_options"]["include_usage"] is True
+
+    assert len(chunks) == len(streaming_tool_call_chunks)
+    assert chunks == streaming_tool_call_chunks
+
+    assert mock_client.capture.call_count == 1
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+
+    assert call_args["distinct_id"] == "test-id"
+    assert call_args["event"] == "$ai_generation"
+    assert props["$ai_provider"] == "openai"
+    assert props["$ai_model"] == "gpt-4"
+    assert props["$ai_output_tokens"] == 15
+    assert props["$ai_input_tokens"] == 20
+    assert isinstance(props["$ai_latency"], float)
+
+
+@pytest.mark.asyncio
+async def test_async_responses_streaming_with_tokens(mock_client):
+    from openai.types.responses import ResponseUsage
+    from unittest.mock import MagicMock
+
+    chunks = []
+
+    chunk1 = MagicMock()
+    chunk1.type = "response.text.delta"
+    chunk1.text = "Test "
+    chunks.append(chunk1)
+
+    chunk2 = MagicMock()
+    chunk2.type = "response.text.delta"
+    chunk2.text = "response"
+    chunks.append(chunk2)
+
+    chunk3 = MagicMock()
+    chunk3.type = "response.completed"
+    chunk3.response = MagicMock()
+    chunk3.response.usage = ResponseUsage(
+        input_tokens=25,
+        output_tokens=30,
+        total_tokens=55,
+        input_tokens_details={"prompt_tokens": 25, "cached_tokens": 0},
+        output_tokens_details={"reasoning_tokens": 0},
+    )
+    chunk3.response.output = ["Test response"]
+    chunks.append(chunk3)
+
+    captured_kwargs = {}
+
+    async def mock_create(self, **kwargs):
+        captured_kwargs["kwargs"] = kwargs
+
+        async def chunk_iterable():
+            for chunk in chunks:
+                yield chunk
+
+        return chunk_iterable()
+
+    with patch("openai.resources.responses.AsyncResponses.create", new=mock_create):
+        client = AsyncOpenAI(api_key="test-key", posthog_client=mock_client)
+
+        response_stream = await client.responses.create(
+            model="gpt-4o-mini",
+            input=[{"role": "user", "content": "Test message"}],
+            stream=True,
+            posthog_distinct_id="test-id",
+            posthog_properties={"test": "streaming"},
+        )
+
+        async for _ in response_stream:
+            pass
+
+    kwargs = captured_kwargs["kwargs"]
+    assert "stream_options" not in kwargs
+
+    assert mock_client.capture.call_count == 1
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+
+    assert call_args["distinct_id"] == "test-id"
+    assert call_args["event"] == "$ai_generation"
+    assert props["$ai_provider"] == "openai"
+    assert props["$ai_model"] == "gpt-4o-mini"
+    assert props["$ai_input_tokens"] == 25
+    assert props["$ai_output_tokens"] == 30
+    assert props["test"] == "streaming"
+    assert isinstance(props["$ai_latency"], float)
+
+
+@pytest.mark.asyncio
+async def test_async_embeddings_create(mock_client, mock_embedding_response):
+    mock_create = AsyncMock(return_value=mock_embedding_response)
+
+    with patch("openai.resources.embeddings.AsyncEmbeddings.create", new=mock_create):
+        client = AsyncOpenAI(api_key="test-key", posthog_client=mock_client)
+
+        response = await client.embeddings.create(
+            model="text-embedding-3-small",
+            input="Hello world",
+            posthog_distinct_id="test-id",
+            posthog_properties={"foo": "bar"},
+        )
+
+    assert response == mock_embedding_response
+    assert mock_create.await_count == 1
+    assert mock_client.capture.call_count == 1
+
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+
+    assert call_args["distinct_id"] == "test-id"
+    assert call_args["event"] == "$ai_embedding"
+    assert props["$ai_provider"] == "openai"
+    assert props["$ai_model"] == "text-embedding-3-small"
+    assert props["foo"] == "bar"
+    assert isinstance(props["$ai_latency"], float)
+
+
 def test_tool_definition(mock_client, mock_openai_response):
     """Test that tools defined in the create function are captured in $ai_tools property"""
     with patch(

posthog/test/ai/test_sanitization.py

@@ -0,0 +1,335 @@
+import unittest
+
+from posthog.ai.sanitization import (
+    redact_base64_data_url,
+    sanitize_openai,
+    sanitize_openai_response,
+    sanitize_anthropic,
+    sanitize_gemini,
+    sanitize_langchain,
+    is_base64_data_url,
+    is_raw_base64,
+    REDACTED_IMAGE_PLACEHOLDER,
+)
+
+
+class TestSanitization(unittest.TestCase):
+    def setUp(self):
+        self.sample_base64_image = "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQ..."
+        self.sample_base64_png = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAUA..."
+        self.regular_url = "https://example.com/image.jpg"
+        self.raw_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUl=="
+
+    def test_is_base64_data_url(self):
+        self.assertTrue(is_base64_data_url(self.sample_base64_image))
+        self.assertTrue(is_base64_data_url(self.sample_base64_png))
+        self.assertFalse(is_base64_data_url(self.regular_url))
+        self.assertFalse(is_base64_data_url("regular text"))
+
+    def test_is_raw_base64(self):
+        self.assertTrue(is_raw_base64(self.raw_base64))
+        self.assertFalse(is_raw_base64("short"))
+        self.assertFalse(is_raw_base64(self.regular_url))
+        self.assertFalse(is_raw_base64("/path/to/file"))
+
+    def test_redact_base64_data_url(self):
+        self.assertEqual(
+            redact_base64_data_url(self.sample_base64_image), REDACTED_IMAGE_PLACEHOLDER
+        )
+        self.assertEqual(
+            redact_base64_data_url(self.sample_base64_png), REDACTED_IMAGE_PLACEHOLDER
+        )
+        self.assertEqual(redact_base64_data_url(self.regular_url), self.regular_url)
+        self.assertEqual(redact_base64_data_url(None), None)
+        self.assertEqual(redact_base64_data_url(123), 123)
+
+    def test_sanitize_openai(self):
+        input_data = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "What is in this image?"},
+                    {
+                        "type": "image_url",
+                        "image_url": {
+                            "url": self.sample_base64_image,
+                            "detail": "high",
+                        },
+                    },
+                ],
+            }
+        ]
+
+        result = sanitize_openai(input_data)
+
+        self.assertEqual(result[0]["content"][0]["text"], "What is in this image?")
+        self.assertEqual(
+            result[0]["content"][1]["image_url"]["url"], REDACTED_IMAGE_PLACEHOLDER
+        )
+        self.assertEqual(result[0]["content"][1]["image_url"]["detail"], "high")
+
+    def test_sanitize_openai_preserves_regular_urls(self):
+        input_data = [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": self.regular_url},
+                    }
+                ],
+            }
+        ]
+
+        result = sanitize_openai(input_data)
+        self.assertEqual(result[0]["content"][0]["image_url"]["url"], self.regular_url)
+
+    def test_sanitize_openai_response(self):
+        input_data = [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "input_image",
+                        "image_url": self.sample_base64_image,
+                    }
+                ],
+            }
+        ]
+
+        result = sanitize_openai_response(input_data)
+        self.assertEqual(
+            result[0]["content"][0]["image_url"], REDACTED_IMAGE_PLACEHOLDER
+        )
+
+    def test_sanitize_anthropic(self):
+        input_data = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "What is in this image?"},
+                    {
+                        "type": "image",
+                        "source": {
+                            "type": "base64",
+                            "media_type": "image/jpeg",
+                            "data": "base64data",
+                        },
+                    },
+                ],
+            }
+        ]
+
+        result = sanitize_anthropic(input_data)
+
+        self.assertEqual(result[0]["content"][0]["text"], "What is in this image?")
+        self.assertEqual(
+            result[0]["content"][1]["source"]["data"], REDACTED_IMAGE_PLACEHOLDER
+        )
+        self.assertEqual(result[0]["content"][1]["source"]["type"], "base64")
+        self.assertEqual(result[0]["content"][1]["source"]["media_type"], "image/jpeg")
+
+    def test_sanitize_gemini(self):
+        input_data = [
+            {
+                "parts": [
+                    {"text": "What is in this image?"},
+                    {
+                        "inline_data": {
+                            "mime_type": "image/jpeg",
+                            "data": "base64data",
+                        }
+                    },
+                ]
+            }
+        ]
+
+        result = sanitize_gemini(input_data)
+
+        self.assertEqual(result[0]["parts"][0]["text"], "What is in this image?")
+        self.assertEqual(
+            result[0]["parts"][1]["inline_data"]["data"], REDACTED_IMAGE_PLACEHOLDER
+        )
+        self.assertEqual(
+            result[0]["parts"][1]["inline_data"]["mime_type"], "image/jpeg"
+        )
+
+    def test_sanitize_langchain_openai_style(self):
+        input_data = [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": self.sample_base64_image},
+                    }
+                ],
+            }
+        ]
+
+        result = sanitize_langchain(input_data)
+        self.assertEqual(
+            result[0]["content"][0]["image_url"]["url"], REDACTED_IMAGE_PLACEHOLDER
+        )
+
+    def test_sanitize_langchain_anthropic_style(self):
+        input_data = [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "image",
+                        "source": {"data": "base64data"},
+                    }
+                ],
+            }
+        ]
+
+        result = sanitize_langchain(input_data)
+        self.assertEqual(
+            result[0]["content"][0]["source"]["data"], REDACTED_IMAGE_PLACEHOLDER
+        )
+
+    def test_sanitize_with_data_url_format(self):
+        # Test that data URLs are properly detected and redacted across providers
+        data_url = "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQABAAD"
+
+        # OpenAI format
+        openai_data = [
+            {
+                "role": "user",
+                "content": [{"type": "image_url", "image_url": {"url": data_url}}],
+            }
+        ]
+        result = sanitize_openai(openai_data)
+        self.assertEqual(
+            result[0]["content"][0]["image_url"]["url"], REDACTED_IMAGE_PLACEHOLDER
+        )
+
+        # Anthropic format
+        anthropic_data = [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "image",
+                        "source": {
+                            "type": "base64",
+                            "media_type": "image/jpeg",
+                            "data": data_url,
+                        },
+                    }
+                ],
+            }
+        ]
+        result = sanitize_anthropic(anthropic_data)
+        self.assertEqual(
+            result[0]["content"][0]["source"]["data"], REDACTED_IMAGE_PLACEHOLDER
+        )
+
+        # LangChain format
+        langchain_data = [
+            {"role": "user", "content": [{"type": "image", "data": data_url}]}
+        ]
+        result = sanitize_langchain(langchain_data)
+        self.assertEqual(result[0]["content"][0]["data"], REDACTED_IMAGE_PLACEHOLDER)
+
+    def test_sanitize_with_raw_base64(self):
+        # Test that raw base64 strings (without data URL prefix) are detected
+        raw_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUl=="
+
+        # Test with Anthropic format
+        anthropic_data = [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "image",
+                        "source": {
+                            "type": "base64",
+                            "media_type": "image/png",
+                            "data": raw_base64,
+                        },
+                    }
+                ],
+            }
+        ]
+        result = sanitize_anthropic(anthropic_data)
+        self.assertEqual(
+            result[0]["content"][0]["source"]["data"], REDACTED_IMAGE_PLACEHOLDER
+        )
+
+        # Test with Gemini format
+        gemini_data = [
+            {"parts": [{"inline_data": {"mime_type": "image/png", "data": raw_base64}}]}
+        ]
+        result = sanitize_gemini(gemini_data)
+        self.assertEqual(
+            result[0]["parts"][0]["inline_data"]["data"], REDACTED_IMAGE_PLACEHOLDER
+        )
+
+    def test_sanitize_preserves_regular_content(self):
+        # Ensure non-base64 content is preserved across all providers
+        regular_url = "https://example.com/image.jpg"
+        text_content = "What do you see?"
+
+        # OpenAI
+        openai_data = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": text_content},
+                    {"type": "image_url", "image_url": {"url": regular_url}},
+                ],
+            }
+        ]
+        result = sanitize_openai(openai_data)
+        self.assertEqual(result[0]["content"][0]["text"], text_content)
+        self.assertEqual(result[0]["content"][1]["image_url"]["url"], regular_url)
+
+        # Anthropic
+        anthropic_data = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": text_content},
+                    {"type": "image", "source": {"type": "url", "url": regular_url}},
+                ],
+            }
+        ]
+        result = sanitize_anthropic(anthropic_data)
+        self.assertEqual(result[0]["content"][0]["text"], text_content)
+        # URL-based images should remain unchanged
+        self.assertEqual(result[0]["content"][1]["source"]["url"], regular_url)
+
+    def test_sanitize_handles_non_dict_content(self):
+        input_data = [{"role": "user", "content": "Just text"}]
+
+        result = sanitize_openai(input_data)
+        self.assertEqual(result, input_data)
+
+    def test_sanitize_handles_none_input(self):
+        self.assertIsNone(sanitize_openai(None))
+        self.assertIsNone(sanitize_anthropic(None))
+        self.assertIsNone(sanitize_gemini(None))
+        self.assertIsNone(sanitize_langchain(None))
+
+    def test_sanitize_handles_single_message(self):
+        input_data = {
+            "role": "user",
+            "content": [
+                {
+                    "type": "image_url",
+                    "image_url": {"url": self.sample_base64_image},
+                }
+            ],
+        }
+
+        result = sanitize_openai(input_data)
+        self.assertEqual(
+            result["content"][0]["image_url"]["url"], REDACTED_IMAGE_PLACEHOLDER
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

posthog/test/ai/test_system_prompts.py

@@ -0,0 +1,354 @@
+"""
+Tests for system prompt capture across all LLM providers.
+
+This test suite ensures that system prompts are correctly captured in analytics
+regardless of how they're passed to the providers:
+- As first message in messages/contents array (standard format)
+- As separate system parameter (Anthropic, OpenAI)
+- As instructions parameter (OpenAI Responses API)
+- As system_instruction parameter (Gemini)
+"""
+
+import time
+import unittest
+from unittest.mock import patch, MagicMock
+
+
+class TestSystemPromptCapture(unittest.TestCase):
+    """Test system prompt capture for all providers."""
+
+    def setUp(self):
+        super().setUp()
+        self.test_system_prompt = "You are a helpful AI assistant."
+        self.test_user_message = "Hello, how are you?"
+        self.test_response = "I'm doing well, thank you!"
+
+        # Create mock PostHog client
+        self.client = MagicMock()
+        self.client.privacy_mode = False
+
+    def _assert_system_prompt_captured(self, captured_input):
+        """Helper to assert system prompt is correctly captured."""
+        self.assertEqual(
+            len(captured_input), 2, "Should have 2 messages (system + user)"
+        )
+        self.assertEqual(
+            captured_input[0]["role"], "system", "First message should be system"
+        )
+        self.assertEqual(
+            captured_input[0]["content"],
+            self.test_system_prompt,
+            "System content should match",
+        )
+        self.assertEqual(
+            captured_input[1]["role"], "user", "Second message should be user"
+        )
+        self.assertEqual(
+            captured_input[1]["content"],
+            self.test_user_message,
+            "User content should match",
+        )
+
+    # OpenAI Tests
+    def test_openai_messages_array_system_prompt(self):
+        """Test OpenAI with system prompt in messages array."""
+        try:
+            from posthog.ai.openai import OpenAI
+            from openai.types.chat import ChatCompletion, ChatCompletionMessage
+            from openai.types.chat.chat_completion import Choice
+            from openai.types.completion_usage import CompletionUsage
+        except ImportError:
+            self.skipTest("OpenAI package not available")
+
+        mock_response = ChatCompletion(
+            id="test",
+            model="gpt-4",
+            object="chat.completion",
+            created=int(time.time()),
+            choices=[
+                Choice(
+                    finish_reason="stop",
+                    index=0,
+                    message=ChatCompletionMessage(
+                        content=self.test_response, role="assistant"
+                    ),
+                )
+            ],
+            usage=CompletionUsage(
+                completion_tokens=10, prompt_tokens=20, total_tokens=30
+            ),
+        )
+
+        with patch(
+            "openai.resources.chat.completions.Completions.create",
+            return_value=mock_response,
+        ):
+            client = OpenAI(posthog_client=self.client, api_key="test")
+
+            messages = [
+                {"role": "system", "content": self.test_system_prompt},
+                {"role": "user", "content": self.test_user_message},
+            ]
+
+            client.chat.completions.create(
+                model="gpt-4", messages=messages, posthog_distinct_id="test-user"
+            )
+
+            self.assertEqual(len(self.client.capture.call_args_list), 1)
+            properties = self.client.capture.call_args_list[0][1]["properties"]
+            self._assert_system_prompt_captured(properties["$ai_input"])
+
+    def test_openai_separate_system_parameter(self):
+        """Test OpenAI with system prompt as separate parameter."""
+        try:
+            from posthog.ai.openai import OpenAI
+            from openai.types.chat import ChatCompletion, ChatCompletionMessage
+            from openai.types.chat.chat_completion import Choice
+            from openai.types.completion_usage import CompletionUsage
+        except ImportError:
+            self.skipTest("OpenAI package not available")
+
+        mock_response = ChatCompletion(
+            id="test",
+            model="gpt-4",
+            object="chat.completion",
+            created=int(time.time()),
+            choices=[
+                Choice(
+                    finish_reason="stop",
+                    index=0,
+                    message=ChatCompletionMessage(
+                        content=self.test_response, role="assistant"
+                    ),
+                )
+            ],
+            usage=CompletionUsage(
+                completion_tokens=10, prompt_tokens=20, total_tokens=30
+            ),
+        )
+
+        with patch(
+            "openai.resources.chat.completions.Completions.create",
+            return_value=mock_response,
+        ):
+            client = OpenAI(posthog_client=self.client, api_key="test")
+
+            messages = [{"role": "user", "content": self.test_user_message}]
+
+            client.chat.completions.create(
+                model="gpt-4",
+                messages=messages,
+                system=self.test_system_prompt,
+                posthog_distinct_id="test-user",
+            )
+
+            self.assertEqual(len(self.client.capture.call_args_list), 1)
+            properties = self.client.capture.call_args_list[0][1]["properties"]
+            self._assert_system_prompt_captured(properties["$ai_input"])
+
+    def test_openai_streaming_system_parameter(self):
+        """Test OpenAI streaming with system parameter."""
+        try:
+            from posthog.ai.openai import OpenAI
+            from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
+            from openai.types.chat.chat_completion_chunk import Choice as ChoiceChunk
+            from openai.types.chat.chat_completion_chunk import ChoiceDelta
+            from openai.types.completion_usage import CompletionUsage
+        except ImportError:
+            self.skipTest("OpenAI package not available")
+
+        chunk1 = ChatCompletionChunk(
+            id="test",
+            model="gpt-4",
+            object="chat.completion.chunk",
+            created=int(time.time()),
+            choices=[
+                ChoiceChunk(
+                    finish_reason=None,
+                    index=0,
+                    delta=ChoiceDelta(content="Hello", role="assistant"),
+                )
+            ],
+        )
+
+        chunk2 = ChatCompletionChunk(
+            id="test",
+            model="gpt-4",
+            object="chat.completion.chunk",
+            created=int(time.time()),
+            choices=[
+                ChoiceChunk(
+                    finish_reason="stop",
+                    index=0,
+                    delta=ChoiceDelta(content=" there!", role=None),
+                )
+            ],
+            usage=CompletionUsage(
+                completion_tokens=10, prompt_tokens=20, total_tokens=30
+            ),
+        )
+
+        with patch(
+            "openai.resources.chat.completions.Completions.create",
+            return_value=[chunk1, chunk2],
+        ):
+            client = OpenAI(posthog_client=self.client, api_key="test")
+
+            messages = [{"role": "user", "content": self.test_user_message}]
+
+            response_generator = client.chat.completions.create(
+                model="gpt-4",
+                messages=messages,
+                system=self.test_system_prompt,
+                stream=True,
+                posthog_distinct_id="test-user",
+            )
+
+            list(response_generator)  # Consume generator
+
+            self.assertEqual(len(self.client.capture.call_args_list), 1)
+            properties = self.client.capture.call_args_list[0][1]["properties"]
+            self._assert_system_prompt_captured(properties["$ai_input"])
+
+    # Anthropic Tests
+    def test_anthropic_messages_array_system_prompt(self):
+        """Test Anthropic with system prompt in messages array."""
+        try:
+            from posthog.ai.anthropic import Anthropic
+        except ImportError:
+            self.skipTest("Anthropic package not available")
+
+        with patch("anthropic.resources.messages.Messages.create") as mock_create:
+            mock_response = MagicMock()
+            mock_response.usage.input_tokens = 20
+            mock_response.usage.output_tokens = 10
+            mock_response.usage.cache_read_input_tokens = None
+            mock_response.usage.cache_creation_input_tokens = None
+            mock_create.return_value = mock_response
+
+            client = Anthropic(posthog_client=self.client, api_key="test")
+
+            messages = [
+                {"role": "system", "content": self.test_system_prompt},
+                {"role": "user", "content": self.test_user_message},
+            ]
+
+            client.messages.create(
+                model="claude-3-5-sonnet-20241022",
+                messages=messages,
+                posthog_distinct_id="test-user",
+            )
+
+            self.assertEqual(len(self.client.capture.call_args_list), 1)
+            properties = self.client.capture.call_args_list[0][1]["properties"]
+            self._assert_system_prompt_captured(properties["$ai_input"])
+
+    def test_anthropic_separate_system_parameter(self):
+        """Test Anthropic with system prompt as separate parameter."""
+        try:
+            from posthog.ai.anthropic import Anthropic
+        except ImportError:
+            self.skipTest("Anthropic package not available")
+
+        with patch("anthropic.resources.messages.Messages.create") as mock_create:
+            mock_response = MagicMock()
+            mock_response.usage.input_tokens = 20
+            mock_response.usage.output_tokens = 10
+            mock_response.usage.cache_read_input_tokens = None
+            mock_response.usage.cache_creation_input_tokens = None
+            mock_create.return_value = mock_response
+
+            client = Anthropic(posthog_client=self.client, api_key="test")
+
+            messages = [{"role": "user", "content": self.test_user_message}]
+
+            client.messages.create(
+                model="claude-3-5-sonnet-20241022",
+                messages=messages,
+                system=self.test_system_prompt,
+                posthog_distinct_id="test-user",
+            )
+
+            self.assertEqual(len(self.client.capture.call_args_list), 1)
+            properties = self.client.capture.call_args_list[0][1]["properties"]
+            self._assert_system_prompt_captured(properties["$ai_input"])
+
+    # Gemini Tests
+    def test_gemini_contents_array_system_prompt(self):
+        """Test Gemini with system prompt in contents array."""
+        try:
+            from posthog.ai.gemini import Client
+        except ImportError:
+            self.skipTest("Gemini package not available")
+
+        with patch("google.genai.Client") as mock_genai_class:
+            mock_response = MagicMock()
+            mock_response.candidates = [MagicMock()]
+            mock_response.candidates[0].content.parts = [MagicMock()]
+            mock_response.candidates[0].content.parts[0].text = self.test_response
+            mock_response.usage_metadata.prompt_token_count = 20
+            mock_response.usage_metadata.candidates_token_count = 10
+            mock_response.usage_metadata.cached_content_token_count = None
+            mock_response.usage_metadata.thoughts_token_count = None
+
+            mock_client_instance = MagicMock()
+            mock_models_instance = MagicMock()
+            mock_models_instance.generate_content.return_value = mock_response
+            mock_client_instance.models = mock_models_instance
+            mock_genai_class.return_value = mock_client_instance
+
+            client = Client(posthog_client=self.client, api_key="test")
+
+            contents = [
+                {"role": "system", "content": self.test_system_prompt},
+                {"role": "user", "content": self.test_user_message},
+            ]
+
+            client.models.generate_content(
+                model="gemini-2.0-flash",
+                contents=contents,
+                posthog_distinct_id="test-user",
+            )
+
+            self.assertEqual(len(self.client.capture.call_args_list), 1)
+            properties = self.client.capture.call_args_list[0][1]["properties"]
+            self._assert_system_prompt_captured(properties["$ai_input"])
+
+    def test_gemini_system_instruction_parameter(self):
+        """Test Gemini with system_instruction in config parameter."""
+        try:
+            from posthog.ai.gemini import Client
+        except ImportError:
+            self.skipTest("Gemini package not available")
+
+        with patch("google.genai.Client") as mock_genai_class:
+            mock_response = MagicMock()
+            mock_response.candidates = [MagicMock()]
+            mock_response.candidates[0].content.parts = [MagicMock()]
+            mock_response.candidates[0].content.parts[0].text = self.test_response
+            mock_response.usage_metadata.prompt_token_count = 20
+            mock_response.usage_metadata.candidates_token_count = 10
+            mock_response.usage_metadata.cached_content_token_count = None
+            mock_response.usage_metadata.thoughts_token_count = None
+
+            mock_client_instance = MagicMock()
+            mock_models_instance = MagicMock()
+            mock_models_instance.generate_content.return_value = mock_response
+            mock_client_instance.models = mock_models_instance
+            mock_genai_class.return_value = mock_client_instance
+
+            client = Client(posthog_client=self.client, api_key="test")
+
+            contents = [{"role": "user", "content": self.test_user_message}]
+            config = {"system_instruction": self.test_system_prompt}
+
+            client.models.generate_content(
+                model="gemini-2.0-flash",
+                contents=contents,
+                config=config,
+                posthog_distinct_id="test-user",
+            )
+
+            self.assertEqual(len(self.client.capture.call_args_list), 1)
+            properties = self.client.capture.call_args_list[0][1]["properties"]
+            self._assert_system_prompt_captured(properties["$ai_input"])

posthog/test/integrations/test_middleware.py

@@ -4,7 +4,21 @@ from posthog.contexts import (
     get_context_distinct_id,
 )
 import unittest
-from unittest.mock import Mock
+from unittest.mock import Mock, patch
+import asyncio
+
+# Configure Django settings before importing middleware
+import django
+from django.conf import settings
+
+if not settings.configured:
+    settings.configure(
+        DEBUG=True,
+        SECRET_KEY="test-secret-key",
+        INSTALLED_APPS=[],
+        MIDDLEWARE=[],
+    )
+    django.setup()
 
 from posthog.integrations.django import PosthogContextMiddleware
 
@@ -38,14 +52,33 @@ class TestPosthogContextMiddleware(unittest.TestCase):
         request_filter=None,
         tag_map=None,
         capture_exceptions=True,
+        get_response=None,
     ):
-        """Helper to create middleware instance without calling __init__"""
-        middleware = PosthogContextMiddleware.__new__(PosthogContextMiddleware)
-        middleware.get_response = Mock()
-        middleware.extra_tags = extra_tags
-        middleware.request_filter = request_filter
-        middleware.tag_map = tag_map
-        middleware.capture_exceptions = capture_exceptions
+        """Helper to create middleware instance with mock Django settings"""
+        if get_response is None:
+            get_response = Mock()
+
+        with patch("django.conf.settings") as mock_settings:
+            # Configure mock settings
+            mock_settings.POSTHOG_MW_EXTRA_TAGS = extra_tags
+            mock_settings.POSTHOG_MW_REQUEST_FILTER = request_filter
+            mock_settings.POSTHOG_MW_TAG_MAP = tag_map
+            mock_settings.POSTHOG_MW_CAPTURE_EXCEPTIONS = capture_exceptions
+            mock_settings.POSTHOG_MW_CLIENT = None
+
+            # Make hasattr work correctly
+            def mock_hasattr(obj, name):
+                return name in [
+                    "POSTHOG_MW_EXTRA_TAGS",
+                    "POSTHOG_MW_REQUEST_FILTER",
+                    "POSTHOG_MW_TAG_MAP",
+                    "POSTHOG_MW_CAPTURE_EXCEPTIONS",
+                    "POSTHOG_MW_CLIENT",
+                ]
+
+            with patch("builtins.hasattr", side_effect=mock_hasattr):
+                middleware = PosthogContextMiddleware(get_response)
+
         return middleware
 
     def test_extract_tags_basic(self):
@@ -168,6 +201,348 @@ class TestPosthogContextMiddleware(unittest.TestCase):
 
         self.assertEqual(tags["$request_method"], "PATCH")
 
+    def test_process_exception_called_during_view_exception(self):
+        """
+        Unit test verifying process_exception captures exceptions per Django's contract.
+
+        Since this is a library test (no Django runtime), we simulate how Django
+        would invoke our middleware in production:
+        1. Middleware.__call__ creates context with request tags
+        2. View raises exception inside get_response
+        3. Django's BaseHandler catches it, calls process_exception, returns error response
+        4. Exception never propagates to middleware's context manager
+
+        We manually call process_exception to simulate Django's behavior - this is
+        the only way to test the hook without a full Django integration test.
+        """
+        mock_client = Mock()
+        view_exception = ValueError("View raised this error")
+        error_response = Mock(status_code=500)
+
+        def mock_get_response(request):
+            # Simulate Django's exception handling: catches view exception,
+            # calls process_exception hook if it exists, returns error response
+            if hasattr(middleware, "process_exception"):
+                middleware.process_exception(request, view_exception)
+            return error_response
+
+        middleware = self.create_middleware(get_response=mock_get_response)
+        middleware.client = mock_client
+
+        request = MockRequest(
+            headers={"X-POSTHOG-DISTINCT-ID": "test-user"},
+            method="POST",
+            path="/api/endpoint",
+        )
+        response = middleware(request)
+
+        self.assertEqual(response.status_code, 500)
+        mock_client.capture_exception.assert_called_once_with(view_exception)
+
+    def test_process_exception_respects_capture_exceptions_false(self):
+        """Verify process_exception respects capture_exceptions=False setting"""
+        mock_client = Mock()
+        view_exception = ValueError("Should not be captured")
+
+        def mock_get_response(request):
+            if hasattr(middleware, "process_exception"):
+                middleware.process_exception(request, view_exception)
+            return Mock(status_code=500)
+
+        middleware = self.create_middleware(
+            capture_exceptions=False, get_response=mock_get_response
+        )
+        middleware.client = mock_client
+
+        request = MockRequest()
+        middleware(request)
+
+        mock_client.capture_exception.assert_not_called()
+
+    def test_process_exception_respects_request_filter(self):
+        """Verify process_exception respects request_filter setting"""
+        mock_client = Mock()
+        view_exception = ValueError("Should be filtered")
+
+        def mock_get_response(request):
+            if hasattr(middleware, "process_exception"):
+                middleware.process_exception(request, view_exception)
+            return Mock(status_code=500)
+
+        middleware = self.create_middleware(
+            request_filter=lambda req: False,
+            capture_exceptions=True,
+            get_response=mock_get_response,
+        )
+        middleware.client = mock_client
+
+        request = MockRequest()
+        middleware(request)
+
+        mock_client.capture_exception.assert_not_called()
+
+
+class TestPosthogContextMiddlewareSync(unittest.TestCase):
+    """Test synchronous middleware behavior"""
+
+    def test_sync_middleware_call(self):
+        """Test that sync middleware correctly processes requests"""
+        mock_response = Mock()
+        get_response = Mock(return_value=mock_response)
+
+        # Create middleware with sync get_response
+        middleware = PosthogContextMiddleware(get_response)
+
+        # Verify sync mode detected
+        self.assertFalse(middleware._is_coroutine)
+
+        request = MockRequest(
+            headers={"X-POSTHOG-SESSION-ID": "test-session"},
+            method="GET",
+            path="/test",
+        )
+
+        with new_context():
+            response = middleware(request)
+
+            # Verify response returned
+            self.assertEqual(response, mock_response)
+            get_response.assert_called_once_with(request)
+
+    def test_sync_middleware_with_filter(self):
+        """Test sync middleware respects request filter"""
+        mock_response = Mock()
+        get_response = Mock(return_value=mock_response)
+
+        # Create middleware with request filter that filters all requests
+        request_filter = lambda req: False
+        middleware = PosthogContextMiddleware.__new__(PosthogContextMiddleware)
+        middleware.get_response = get_response
+        middleware._is_coroutine = False
+        middleware.request_filter = request_filter
+        middleware.capture_exceptions = True
+        middleware.client = None
+
+        request = MockRequest()
+
+        # Should skip context creation and return response directly
+        response = middleware(request)
+        self.assertEqual(response, mock_response)
+        get_response.assert_called_once_with(request)
+
+    def test_view_exceptions_only_captured_via_process_exception(self):
+        """
+        Demonstrates that process_exception is required to capture view exceptions.
+
+        In production Django, view exceptions don't propagate to middleware's context
+        manager because Django's BaseHandler catches them first and converts them to
+        error responses. Django provides the exception via process_exception hook instead.
+
+        This unit test proves:
+        1. Context manager in __call__ never sees view exceptions (Django intercepts)
+        2. Only process_exception can capture them
+        3. Without process_exception, exceptions are silently lost (v6.7.5 regression)
+
+        We manually call process_exception to verify the hook works - in production,
+        Django's BaseHandler would call it when a view raises.
+        """
+        mock_client = Mock()
+        get_response = Mock(return_value=Mock(status_code=500))
+
+        middleware = PosthogContextMiddleware(get_response)
+        middleware.client = mock_client
+
+        def get_response_simulating_django(request):
+            # Simulates Django behavior: view exception converted to error response,
+            # never propagates to middleware's context manager
+            return Mock(status_code=500)
+
+        middleware._sync_get_response = get_response_simulating_django
+
+        request = MockRequest()
+
+        response = middleware(request)
+        self.assertEqual(response.status_code, 500)
+
+        # Context manager didn't capture anything - exception was intercepted by Django
+        mock_client.capture_exception.assert_not_called()
+
+        # Verify process_exception hook exists and captures exceptions when called
+        if hasattr(middleware, "process_exception"):
+            exception = ValueError("View error")
+            middleware.process_exception(request, exception)
+            mock_client.capture_exception.assert_called_once_with(exception)
+        else:
+            self.fail(
+                "process_exception missing - view exceptions will not be captured!"
+            )
+
+
+class TestPosthogContextMiddlewareAsync(unittest.TestCase):
+    """Test asynchronous middleware behavior"""
+
+    def test_async_middleware_detection(self):
+        """Test that async get_response is correctly detected"""
+
+        async def async_get_response(request):
+            return Mock()
+
+        middleware = PosthogContextMiddleware(async_get_response)
+
+        # Verify async mode detected
+        self.assertTrue(middleware._is_coroutine)
+
+    def test_async_middleware_call(self):
+        """Test that async middleware correctly processes requests"""
+
+        async def run_test():
+            mock_response = Mock()
+
+            async def async_get_response(request):
+                return mock_response
+
+            middleware = PosthogContextMiddleware(async_get_response)
+
+            request = MockRequest(
+                headers={"X-POSTHOG-SESSION-ID": "async-session"},
+                method="POST",
+                path="/async-test",
+            )
+
+            with new_context():
+                # Call should return the coroutine from __acall__
+                result = middleware(request)
+
+                # Verify it's a coroutine
+                self.assertTrue(asyncio.iscoroutine(result))
+
+                # Await the result
+                response = await result
+                self.assertEqual(response, mock_response)
+
+        asyncio.run(run_test())
+
+    def test_async_middleware_with_filter(self):
+        """Test async middleware respects request filter"""
+
+        async def run_test():
+            mock_response = Mock()
+
+            async def async_get_response(request):
+                return mock_response
+
+            # Properly initialize middleware
+            middleware = PosthogContextMiddleware(async_get_response)
+            # Override request filter after initialization
+            middleware.request_filter = lambda req: False
+
+            request = MockRequest()
+
+            # Should skip context creation and return response directly
+            result = middleware(request)
+            response = await result
+            self.assertEqual(response, mock_response)
+
+        asyncio.run(run_test())
+
+    def test_async_middleware_context_propagation(self):
+        """Test that async middleware properly propagates context"""
+
+        async def run_test():
+            mock_response = Mock()
+
+            async def async_get_response(request):
+                # Verify context is available during async processing
+                session_id = get_context_session_id()
+                self.assertEqual(session_id, "async-session-123")
+                return mock_response
+
+            middleware = PosthogContextMiddleware(async_get_response)
+
+            request = MockRequest(
+                headers={"X-POSTHOG-SESSION-ID": "async-session-123"},
+                method="GET",
+            )
+
+            with new_context():
+                result = middleware(request)
+                await result
+
+        asyncio.run(run_test())
+
+    def test_async_middleware_exception_capture(self):
+        """Test that async middleware captures exceptions during request processing"""
+
+        async def run_test():
+            mock_client = Mock()
+
+            # Make async_get_response raise an exception
+            async def raise_exception(request):
+                raise ValueError("Async test exception")
+
+            # Properly initialize middleware
+            middleware = PosthogContextMiddleware(raise_exception)
+            middleware.client = mock_client  # Override with mock client
+
+            request = MockRequest()
+
+            # Should capture exception and re-raise
+            with self.assertRaises(ValueError):
+                result = middleware(request)
+                await result
+
+            # Verify exception was captured by middleware
+            mock_client.capture_exception.assert_called_once()
+            captured_exception = mock_client.capture_exception.call_args[0][0]
+            self.assertIsInstance(captured_exception, ValueError)
+            self.assertEqual(str(captured_exception), "Async test exception")
+
+        asyncio.run(run_test())
+
+
+class TestPosthogContextMiddlewareHybrid(unittest.TestCase):
+    """Test hybrid middleware behavior with mixed sync/async chains"""
+
+    def test_hybrid_flags_set(self):
+        """Test that both capability flags are set"""
+        self.assertTrue(PosthogContextMiddleware.sync_capable)
+        self.assertTrue(PosthogContextMiddleware.async_capable)
+
+    def test_sync_to_async_routing(self):
+        """Test that __call__ routes to __acall__ when async"""
+
+        async def run_test():
+            async def async_get_response(request):
+                return Mock()
+
+            middleware = PosthogContextMiddleware(async_get_response)
+
+            # Verify routing happens
+            request = MockRequest()
+            result = middleware(request)
+
+            # Should be a coroutine from __acall__
+            self.assertTrue(asyncio.iscoroutine(result))
+            await result  # Clean up
+
+        asyncio.run(run_test())
+
+    def test_sync_path_direct_return(self):
+        """Test that sync path returns directly without coroutine"""
+        mock_response = Mock()
+
+        def sync_get_response(request):
+            return mock_response
+
+        middleware = PosthogContextMiddleware(sync_get_response)
+
+        request = MockRequest()
+        result = middleware(request)
+
+        # Should NOT be a coroutine
+        self.assertFalse(asyncio.iscoroutine(result))
+        self.assertEqual(result, mock_response)
+
 
 if __name__ == "__main__":
     unittest.main()

posthog/test/test_client.py

@@ -2423,3 +2423,46 @@ class TestClient(unittest.TestCase):
             batch_data = mock_post.call_args[1]["batch"]
             msg = batch_data[0]
             self.assertEqual(msg["properties"]["$context_tags"], ["random_tag"])
+
+    @mock.patch(
+        "posthog.client.Client._enqueue", side_effect=Exception("Unexpected error")
+    )
+    def test_methods_handle_exceptions(self, mock_enqueue):
+        """Test that all decorated methods handle exceptions gracefully."""
+        client = Client("test-key")
+
+        test_cases = [
+            ("capture", ["test_event"], {}),
+            ("set", [], {"distinct_id": "some-id", "properties": {"a": "b"}}),
+            ("set_once", [], {"distinct_id": "some-id", "properties": {"a": "b"}}),
+            ("group_identify", ["group-type", "group-key"], {}),
+            ("alias", ["some-id", "new-id"], {}),
+        ]
+
+        for method_name, args, kwargs in test_cases:
+            with self.subTest(method=method_name):
+                method = getattr(client, method_name)
+                result = method(*args, **kwargs)
+                self.assertEqual(result, None)
+
+    @mock.patch(
+        "posthog.client.Client._enqueue", side_effect=Exception("Expected error")
+    )
+    def test_debug_flag_re_raises_exceptions(self, mock_enqueue):
+        """Test that methods re-raise exceptions when debug=True."""
+        client = Client("test-key", debug=True)
+
+        test_cases = [
+            ("capture", ["test_event"], {}),
+            ("set", [], {"distinct_id": "some-id", "properties": {"a": "b"}}),
+            ("set_once", [], {"distinct_id": "some-id", "properties": {"a": "b"}}),
+            ("group_identify", ["group-type", "group-key"], {}),
+            ("alias", ["some-id", "new-id"], {}),
+        ]
+
+        for method_name, args, kwargs in test_cases:
+            with self.subTest(method=method_name):
+                method = getattr(client, method_name)
+                with self.assertRaises(Exception) as cm:
+                    method(*args, **kwargs)
+                self.assertEqual(str(cm.exception), "Expected error")

posthog/test/test_feature_flags.py

@@ -365,7 +365,7 @@ class TestLocalEvaluation(unittest.TestCase):
 
     @mock.patch("posthog.client.flags")
     @mock.patch("posthog.client.get")
-    def test_feature_flags_fallback_to_decide(self, patch_get, patch_flags):
+    def test_feature_flags_fallback_to_flags(self, patch_get, patch_flags):
         patch_flags.return_value = {
             "featureFlags": {"beta-feature": "alakazam", "beta-feature2": "alakazam2"}
         }
@@ -431,7 +431,7 @@ class TestLocalEvaluation(unittest.TestCase):
 
     @mock.patch("posthog.client.flags")
     @mock.patch("posthog.client.get")
-    def test_feature_flags_dont_fallback_to_decide_when_only_local_evaluation_is_true(
+    def test_feature_flags_dont_fallback_to_flags_when_only_local_evaluation_is_true(
         self, patch_get, patch_flags
     ):
         patch_flags.return_value = {
@@ -2804,73 +2804,61 @@ class TestLocalEvaluation(unittest.TestCase):
         self.assertEqual(patch_flags.call_count, 0)
 
     @mock.patch("posthog.client.flags")
-    def test_flag_with_multiple_variant_overrides(self, patch_flags):
-        patch_flags.return_value = {"featureFlags": {"beta-feature": "variant-1"}}
+    def test_conditions_evaluated_in_order(self, patch_flags):
+        patch_flags.return_value = {"featureFlags": {"order-test": "server-variant"}}
         client = Client(FAKE_TEST_API_KEY, personal_api_key="test")
         client.feature_flags = [
             {
                 "id": 1,
-                "name": "Beta Feature",
-                "key": "beta-feature",
+                "name": "Order Test Flag",
+                "key": "order-test",
                 "active": True,
-                "rollout_percentage": 100,
                 "filters": {
                     "groups": [
                         {
                             "rollout_percentage": 100,
-                            # The override applies even if the first condition matches all and gives everyone their default group
                         },
                         {
                             "properties": [
                                 {
                                     "key": "email",
                                     "type": "person",
-                                    "value": "test@posthog.com",
-                                    "operator": "exact",
+                                    "value": "@vip.com",
+                                    "operator": "icontains",
                                 }
                             ],
                             "rollout_percentage": 100,
-                            "variant": "second-variant",
+                            "variant": "vip-variant",
                         },
-                        {"rollout_percentage": 50, "variant": "third-variant"},
                     ],
                     "multivariate": {
                         "variants": [
                             {
-                                "key": "first-variant",
-                                "name": "First Variant",
-                                "rollout_percentage": 50,
+                                "key": "control",
+                                "name": "Control",
+                                "rollout_percentage": 100,
                             },
                             {
-                                "key": "second-variant",
-                                "name": "Second Variant",
-                                "rollout_percentage": 25,
-                            },
-                            {
-                                "key": "third-variant",
-                                "name": "Third Variant",
-                                "rollout_percentage": 25,
+                                "key": "vip-variant",
+                                "name": "VIP Variant",
+                                "rollout_percentage": 0,
                             },
                         ]
                     },
                 },
             }
         ]
-        self.assertEqual(
-            client.get_feature_flag(
-                "beta-feature",
-                "test_id",
-                person_properties={"email": "test@posthog.com"},
-            ),
-            "second-variant",
+
+        # Even though user@vip.com would match the second condition with variant override,
+        # they should match the first condition and get control
+        result = client.get_feature_flag(
+            "order-test",
+            "user123",
+            person_properties={"email": "user@vip.com"},
         )
-        self.assertEqual(
-            client.get_feature_flag("beta-feature", "example_id"), "third-variant"
-        )
-        self.assertEqual(
-            client.get_feature_flag("beta-feature", "another_id"), "second-variant"
-        )
-        # decide not called because this can be evaluated locally
+        self.assertEqual(result, "control")
+
+        # server not called because this can be evaluated locally
         self.assertEqual(patch_flags.call_count, 0)
 
     @mock.patch("posthog.client.flags")
@@ -3025,6 +3013,75 @@ class TestLocalEvaluation(unittest.TestCase):
         )
         self.assertEqual(patch_flags.call_count, 0)
 
+    @mock.patch("posthog.client.flags")
+    @mock.patch("posthog.client.get")
+    def test_fallback_to_api_when_flag_has_static_cohort_in_multi_condition(
+        self, patch_get, patch_flags
+    ):
+        """
+        When a flag has multiple conditions and one contains a static cohort,
+        the SDK should fallback to API for the entire flag, not just skip that
+        condition and evaluate the next one locally.
+
+        This prevents returning wrong variants when later conditions could match
+        locally but the user is actually in the static cohort.
+        """
+        client = Client(FAKE_TEST_API_KEY, personal_api_key=FAKE_TEST_API_KEY)
+
+        # Mock the local flags response - cohort 999 is NOT in cohorts map (static cohort)
+        client.feature_flags = [
+            {
+                "id": 1,
+                "key": "multi-condition-flag",
+                "active": True,
+                "filters": {
+                    "groups": [
+                        {
+                            "properties": [
+                                {"key": "id", "value": 999, "type": "cohort"}
+                            ],
+                            "rollout_percentage": 100,
+                            "variant": "set-1",
+                        },
+                        {
+                            "properties": [
+                                {
+                                    "key": "$geoip_country_code",
+                                    "operator": "exact",
+                                    "value": ["DE"],
+                                    "type": "person",
+                                }
+                            ],
+                            "rollout_percentage": 100,
+                            "variant": "set-8",
+                        },
+                    ],
+                    "multivariate": {
+                        "variants": [
+                            {"key": "set-1", "rollout_percentage": 50},
+                            {"key": "set-8", "rollout_percentage": 50},
+                        ]
+                    },
+                },
+            }
+        ]
+        client.cohorts = {}  # Note: cohort 999 is NOT here - it's a static cohort
+
+        # Mock the API response - user is in the static cohort
+        patch_flags.return_value = {"featureFlags": {"multi-condition-flag": "set-1"}}
+
+        result = client.get_feature_flag(
+            "multi-condition-flag",
+            "test-distinct-id",
+            person_properties={"$geoip_country_code": "DE"},
+        )
+
+        # Should return the API result (set-1), not local evaluation (set-8)
+        self.assertEqual(result, "set-1")
+
+        # Verify API was called (fallback occurred)
+        self.assertEqual(patch_flags.call_count, 1)
+
 
 class TestMatchProperties(unittest.TestCase):
     def property(self, key, value, operator=None):
@@ -4018,6 +4075,60 @@ class TestCaptureCalls(unittest.TestCase):
 
         patch_capture.reset_mock()
 
+    @mock.patch("posthog.client.flags")
+    def test_fallback_to_api_in_get_feature_flag_payload_when_flag_has_static_cohort(
+        self, patch_flags
+    ):
+        """
+        Test that get_feature_flag_payload falls back to API when evaluating
+        a flag with static cohorts, similar to get_feature_flag behavior.
+        """
+        client = Client(FAKE_TEST_API_KEY, personal_api_key=FAKE_TEST_API_KEY)
+
+        # Mock the local flags response - cohort 999 is NOT in cohorts map (static cohort)
+        client.feature_flags = [
+            {
+                "id": 1,
+                "name": "Multi-condition Flag",
+                "key": "multi-condition-flag",
+                "active": True,
+                "filters": {
+                    "groups": [
+                        {
+                            "properties": [
+                                {"key": "id", "value": 999, "type": "cohort"}
+                            ],
+                            "rollout_percentage": 100,
+                            "variant": "variant-1",
+                        }
+                    ],
+                    "multivariate": {
+                        "variants": [{"key": "variant-1", "rollout_percentage": 100}]
+                    },
+                    "payloads": {"variant-1": '{"message": "local-payload"}'},
+                },
+            }
+        ]
+        client.cohorts = {}  # Note: cohort 999 is NOT here - it's a static cohort
+
+        # Mock the API response - user is in the static cohort
+        patch_flags.return_value = {
+            "featureFlags": {"multi-condition-flag": "variant-1"},
+            "featureFlagPayloads": {"multi-condition-flag": '{"message": "from-api"}'},
+        }
+
+        # Call get_feature_flag_payload without match_value to trigger evaluation
+        result = client.get_feature_flag_payload(
+            "multi-condition-flag",
+            "test-distinct-id",
+        )
+
+        # Should return the API payload, not local payload
+        self.assertEqual(result, {"message": "from-api"})
+
+        # Verify API was called (fallback occurred)
+        self.assertEqual(patch_flags.call_count, 1)
+
     @mock.patch.object(Client, "capture")
     @mock.patch("posthog.client.flags")
     def test_disable_geoip_get_flag_capture_call(self, patch_flags, patch_capture):

posthog/test/test_module.py

@@ -18,14 +18,6 @@ class TestModule(unittest.TestCase):
             "testsecret", host="http://localhost:8000", on_error=self.failed
         )
 
-    def test_no_api_key(self):
-        self.posthog.api_key = None
-        self.assertRaises(Exception, self.posthog.capture)
-
-    def test_no_host(self):
-        self.posthog.host = None
-        self.assertRaises(Exception, self.posthog.capture)
-
     def test_track(self):
         res = self.posthog.capture("python module event", distinct_id="distinct_id")
         self._assert_enqueue_result(res)

posthog/version.py

@@ -1,4 +1,4 @@
-VERSION = "6.7.0"
+VERSION = "6.7.11"
 
 if __name__ == "__main__":
     print(VERSION, end="")  # noqa: T201

test_capture_ai_local.py

@@ -0,0 +1,326 @@
+#!/usr/bin/env python3
+"""
+Test script to send capture_ai events to localhost:8010.
+This script tests the actual network request to a local PostHog instance.
+"""
+
+from posthog import Posthog
+from uuid import uuid4
+
+# Create a client pointing to localhost:8010
+posthog = Posthog(
+    "test-api-key",  # Use your actual project API key if needed
+    host="http://localhost:8010",
+    debug=True,  # Enable debug mode to see detailed logs
+)
+
+print("Testing capture_ai with localhost:8010")
+print("=" * 60)
+
+# Test 1: $ai_generation event with blobs
+print("\n1. Testing $ai_generation event with blobs...")
+print("-" * 60)
+
+trace_id = f"trace_{uuid4().hex[:8]}"
+
+try:
+    event_uuid = posthog.capture_ai(
+        "$ai_generation",
+        distinct_id="test_user_123",
+        properties={
+            "$ai_model": "gpt-4",
+            "$ai_provider": "openai",
+            "$ai_trace_id": trace_id,
+            "$ai_input": {
+                "messages": [
+                    {
+                        "role": "system",
+                        "content": "You are a helpful assistant that answers questions about Python.",
+                    },
+                    {
+                        "role": "user",
+                        "content": "What is the difference between a list and a tuple?",
+                    },
+                ],
+                "temperature": 0.7,
+                "max_tokens": 500,
+            },
+            "$ai_output_choices": {
+                "choices": [
+                    {
+                        "index": 0,
+                        "message": {
+                            "role": "assistant",
+                            "content": "A list is mutable (can be changed) while a tuple is immutable (cannot be changed after creation). Lists use square brackets [] and tuples use parentheses ().",
+                        },
+                        "finish_reason": "stop",
+                    }
+                ],
+                "model": "gpt-4",
+                "usage": {
+                    "prompt_tokens": 45,
+                    "completion_tokens": 32,
+                    "total_tokens": 77,
+                },
+            },
+            "$ai_completion_tokens": 32,
+            "$ai_prompt_tokens": 45,
+            "$ai_total_tokens": 77,
+            "$ai_latency": 1.234,
+        },
+        blob_properties=["$ai_input", "$ai_output_choices"],
+    )
+
+    if event_uuid:
+        print("✓ SUCCESS: $ai_generation event sent")
+        print(f"  UUID: {event_uuid}")
+        print(f"  Trace ID: {trace_id}")
+    else:
+        print("✗ FAILED: No UUID returned")
+
+except Exception as e:
+    print(f"✗ ERROR: {e}")
+    import traceback
+
+    traceback.print_exc()
+
+# Test 2: $ai_trace event
+print("\n2. Testing $ai_trace event...")
+print("-" * 60)
+
+try:
+    event_uuid = posthog.capture_ai(
+        "$ai_trace",
+        distinct_id="test_user_123",
+        properties={
+            "$ai_model": "gpt-4",
+            "$ai_trace_id": trace_id,
+            "$ai_trace_name": "python_qa_session",
+            "$ai_input_state": {
+                "session_id": "session_123",
+                "user_context": "learning Python",
+            },
+            "$ai_output_state": {"questions_answered": 1, "satisfaction_score": 5},
+        },
+        blob_properties=["$ai_input_state", "$ai_output_state"],
+    )
+
+    if event_uuid:
+        print("✓ SUCCESS: $ai_trace event sent")
+        print(f"  UUID: {event_uuid}")
+        print(f"  Trace ID: {trace_id}")
+    else:
+        print("✗ FAILED: No UUID returned")
+
+except Exception as e:
+    print(f"✗ ERROR: {e}")
+    import traceback
+
+    traceback.print_exc()
+
+# Test 3: $ai_span event
+print("\n3. Testing $ai_span event...")
+print("-" * 60)
+
+span_id = f"span_{uuid4().hex[:8]}"
+
+try:
+    event_uuid = posthog.capture_ai(
+        "$ai_span",
+        distinct_id="test_user_123",
+        properties={
+            "$ai_model": "gpt-4",
+            "$ai_trace_id": trace_id,
+            "$ai_span_id": span_id,
+            "$ai_span_name": "answer_generation",
+            "$ai_parent_id": trace_id,
+            "$ai_span_kind": "llm",
+            "$ai_latency": 0.8,
+        },
+    )
+
+    if event_uuid:
+        print("✓ SUCCESS: $ai_span event sent")
+        print(f"  UUID: {event_uuid}")
+        print(f"  Span ID: {span_id}")
+    else:
+        print("✗ FAILED: No UUID returned")
+
+except Exception as e:
+    print(f"✗ ERROR: {e}")
+    import traceback
+
+    traceback.print_exc()
+
+# Test 4: $ai_embedding event
+print("\n4. Testing $ai_embedding event...")
+print("-" * 60)
+
+try:
+    event_uuid = posthog.capture_ai(
+        "$ai_embedding",
+        distinct_id="test_user_123",
+        properties={
+            "$ai_model": "text-embedding-ada-002",
+            "$ai_provider": "openai",
+            "$ai_trace_id": trace_id,
+            "$ai_input": {
+                "text": "What is the difference between a list and a tuple in Python?"
+            },
+            "$ai_embedding_dimension": 1536,
+            "$ai_latency": 0.123,
+        },
+        blob_properties=["$ai_input"],
+    )
+
+    if event_uuid:
+        print("✓ SUCCESS: $ai_embedding event sent")
+        print(f"  UUID: {event_uuid}")
+    else:
+        print("✗ FAILED: No UUID returned")
+
+except Exception as e:
+    print(f"✗ ERROR: {e}")
+    import traceback
+
+    traceback.print_exc()
+
+# Test 5: $ai_metric event
+print("\n5. Testing $ai_metric event...")
+print("-" * 60)
+
+try:
+    event_uuid = posthog.capture_ai(
+        "$ai_metric",
+        distinct_id="test_user_123",
+        properties={
+            "$ai_model": "gpt-4",
+            "$ai_trace_id": trace_id,
+            "$ai_metric_name": "response_quality",
+            "$ai_metric_value": "0.95",
+        },
+    )
+
+    if event_uuid:
+        print("✓ SUCCESS: $ai_metric event sent")
+        print(f"  UUID: {event_uuid}")
+    else:
+        print("✗ FAILED: No UUID returned")
+
+except Exception as e:
+    print(f"✗ ERROR: {e}")
+    import traceback
+
+    traceback.print_exc()
+
+# Test 6: $ai_feedback event
+print("\n6. Testing $ai_feedback event...")
+print("-" * 60)
+
+try:
+    event_uuid = posthog.capture_ai(
+        "$ai_feedback",
+        distinct_id="test_user_123",
+        properties={
+            "$ai_model": "gpt-4",
+            "$ai_trace_id": trace_id,
+            "$ai_feedback_text": "Great explanation! Very clear and helpful.",
+            "$ai_feedback_rating": 5,
+        },
+    )
+
+    if event_uuid:
+        print("✓ SUCCESS: $ai_feedback event sent")
+        print(f"  UUID: {event_uuid}")
+    else:
+        print("✗ FAILED: No UUID returned")
+
+except Exception as e:
+    print(f"✗ ERROR: {e}")
+    import traceback
+
+    traceback.print_exc()
+
+# Test 7: Test with custom blob properties
+print("\n7. Testing with custom blob properties...")
+print("-" * 60)
+
+try:
+    event_uuid = posthog.capture_ai(
+        "$ai_generation",
+        distinct_id="test_user_123",
+        properties={
+            "$ai_model": "claude-3-opus",
+            "$ai_provider": "anthropic",
+            "$ai_trace_id": trace_id,
+            "$ai_input": {
+                "messages": [{"role": "user", "content": "Write a haiku about Python"}]
+            },
+            "$ai_output_choices": {
+                "choices": [
+                    {
+                        "message": {
+                            "role": "assistant",
+                            "content": "Snake glides through code\nSimple syntax, powerful tools\nDevelopers smile",
+                        }
+                    }
+                ]
+            },
+            "$ai_custom_data": {
+                "large_context": "This is some large custom data that should be sent as a blob"
+            },
+            "$ai_completion_tokens": 20,
+            "$ai_prompt_tokens": 10,
+        },
+        # Custom blob properties - including the default ones plus a custom one
+        blob_properties=["$ai_input", "$ai_output_choices", "$ai_custom_data"],
+    )
+
+    if event_uuid:
+        print("✓ SUCCESS: Event with custom blob properties sent")
+        print(f"  UUID: {event_uuid}")
+    else:
+        print("✗ FAILED: No UUID returned")
+
+except Exception as e:
+    print(f"✗ ERROR: {e}")
+    import traceback
+
+    traceback.print_exc()
+
+# Test 8: Test with groups
+print("\n8. Testing with groups...")
+print("-" * 60)
+
+try:
+    event_uuid = posthog.capture_ai(
+        "$ai_generation",
+        distinct_id="test_user_123",
+        groups={"company": "posthog_inc", "team": "engineering"},
+        properties={
+            "$ai_model": "gpt-4",
+            "$ai_provider": "openai",
+            "$ai_trace_id": trace_id,
+            "$ai_input": {"messages": [{"role": "user", "content": "test"}]},
+            "$ai_output_choices": {
+                "choices": [{"message": {"role": "assistant", "content": "response"}}]
+            },
+        },
+    )
+
+    if event_uuid:
+        print("✓ SUCCESS: Event with groups sent")
+        print(f"  UUID: {event_uuid}")
+    else:
+        print("✗ FAILED: No UUID returned")
+
+except Exception as e:
+    print(f"✗ ERROR: {e}")
+    import traceback
+
+    traceback.print_exc()
+
+print("\n" + "=" * 60)
+print("All tests completed!")
+print("\nMake sure your local PostHog instance is running on http://localhost:8010")
+print("and that the /i/v0/ai endpoint is available.")