fix(llma): missing await in OpenAI's streaming implementation (#342)

* fix(llma): missing async for OpenAI async

* chore(llma): bump version

* chore(llma): bump version

.env.example

@@ -0,0 +1,11 @@
+# PostHog API Configuration
+# Copy this file to .env and update with your actual values
+
+# Your project API key (found on the /setup page in PostHog)
+POSTHOG_PROJECT_API_KEY=phc_your_project_api_key_here
+
+# Your personal API key (for local evaluation and other advanced features)
+POSTHOG_PERSONAL_API_KEY=phx_your_personal_api_key_here
+
+# PostHog host URL (remove this line if using posthog.com)
+POSTHOG_HOST=http://localhost:8000

.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.yaml

@@ -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.yaml

@@ -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,6 +1,64 @@
+# 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
+
+# 6.6.1 - 2025-08-21
+
+- fix: Prevent `NoneType` error when `group_properties` is `None`
+
+# 6.6.0 - 2025-08-15
+
+- feat: Add `flag_keys_to_evaluate` parameter to optimize feature flag evaluation performance by only evaluating specified flags
+- feat: Add `flag_keys_filter` option to `send_feature_flags` for selective flag evaluation in capture events
+
+# 6.5.0 - 2025-08-08
+
+- feat: Add `$context_tags` to an event to know which properties were included as tags
+
+# 6.4.1 - 2025-08-06
+
+- fix: Always pass project API key in `remote_config` requests for deterministic project routing
+
+# 6.4.0 - 2025-08-05
+
+- feat: support Vertex AI for Gemini
+
 # 6.3.4 - 2025-08-04
 
-- fix: Set `$ai_tools` for all providers and `$ai_output_choices` for all non-streaming provider flows properly
+- fix: set `$ai_tools` for all providers and `$ai_output_choices` for all non-streaming provider flows properly
 
 # 6.3.3 - 2025-08-01
 

bin/docs_scripts/doc_constant.py

@@ -3,50 +3,11 @@ Constants for PostHog Python SDK documentation generation.
 """
 
 from typing import Dict, Union
-
-# Types that are built-in to Python and don't need to be documented
-NO_DOCS_TYPES = [
-    "Client",
-    "any",
-    "int",
-    "float",
-    "bool",
-    "dict",
-    "list",
-    "str",
-    "tuple",
-    "set",
-    "frozenset",
-    "bytes",
-    "bytearray",
-    "memoryview",
-    "range",
-    "slice",
-    "complex",
-    "Union",
-    "Optional",
-    "Any",
-    "Callable",
-    "Type",
-    "TypeVar",
-    "Generic",
-    "Literal",
-    "ClassVar",
-    "Final",
-    "Annotated",
-    "NotRequired",
-    "Required",
-    "None",
-    "NoneType",
-    "object",
-    "Unpack",
-    "BaseException",
-    "Exception",
-]
+from posthog.version import VERSION
 
 # Documentation generation metadata
 DOCUMENTATION_METADATA = {
-    "hogRef": "0.1",
+    "hogRef": "0.3",
     "slugPrefix": "posthog-python",
     "specUrl": "https://github.com/PostHog/posthog-python",
 }
@@ -67,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

@@ -11,7 +11,6 @@ from dataclasses import is_dataclass, fields
 from typing import get_origin, get_args, Union
 from textwrap import dedent
 from doc_constant import (
-    NO_DOCS_TYPES,
     DOCUMENTATION_METADATA,
     DOCSTRING_PATTERNS,
     OUTPUT_CONFIG,
@@ -187,7 +186,7 @@ def analyze_parameter(param: inspect.Parameter, docstring: str = "") -> dict:
         param_type = get_type_name(type(param.default))
 
     # Extract parameter description from Args section
-    param_description = f"Parameter: {param.name}"
+    param_description = ""
     if docstring:
         # Look for Args section and extract description for this parameter
         args_section_match = re.search(
@@ -378,6 +377,14 @@ def generate_sdk_documentation():
             except Exception as e:
                 print(f"Error analyzing type {name}: {e}")
 
+    # Clean types of empty types
+
+    # Remove types that have no properties and no examples
+    # Remove types that have no properties and no examples
+    types_list = [
+        t for t in types_list if len(t["properties"]) > 0 or t["example"] != ""
+    ]
+
     # Collect classes
     classes_list = []
 
@@ -420,14 +427,28 @@ def generate_sdk_documentation():
             }
         )
 
+    # Collect categories from functions
+    categories = ["Initialization", "Identification", "Capture"]
+    seen_categories = set(categories)
+    for class_info in classes_list:
+        if "functions" in class_info:
+            for func in class_info["functions"]:
+                if (
+                    "category" in func
+                    and func["category"] not in seen_categories
+                    and func["category"]
+                ):
+                    categories.append(func["category"])
+                    seen_categories.add(func["category"])
+
     # Create the final structure
     result = {
         "id": "posthog-python",
         "hogRef": DOCUMENTATION_METADATA["hogRef"],
         "info": sdk_info,
-        "noDocsTypes": NO_DOCS_TYPES,
         "types": types_list,
         "classes": classes_list,
+        "categories": categories,
     }
 
     return result
@@ -439,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}")
 
@@ -459,12 +491,6 @@ if __name__ == "__main__":
         print(f"   • {classes_count} classes documented")
         print(f"   • {total_functions} functions documented")
 
-        no_docs = documentation["noDocsTypes"]
-        if no_docs:
-            print(
-                f"   • {len(no_docs)} types without documentation: {', '.join(no_docs[:5])}{'...' if len(no_docs) > 5 else ''}"
-            )
-
     except Exception as e:
         print(f"❌ Error generating documentation: {e}")
         import traceback

bin/fmt

@@ -6,9 +6,7 @@ set_source_and_root_dir
 ensure_virtual_env
 
 if [[ "$1" == "--check" ]]; then
-  black --check .
-  isort --check-only .
+  ruff format --check .
 else
-  black .
-  isort .
+  ruff format .
 fi

example.py

@@ -1,186 +1,502 @@
 # PostHog Python library example
-import argparse
+#
+# This script demonstrates various PostHog Python SDK capabilities including:
+# - Basic event capture and user identification
+# - Feature flag local evaluation
+# - Feature flag payloads
+# - Context management and tagging
+#
+# Setup:
+# 1. Copy .env.example to .env and fill in your PostHog credentials
+# 2. Run this script and choose from the interactive menu
+
+import os
 
 import posthog
 
-# Add argument parsing
-parser = argparse.ArgumentParser(description="PostHog Python library example")
-parser.add_argument(
-    "--flag",
-    default="person-on-events-enabled",
-    help="Feature flag key to check (default: person-on-events-enabled)",
-)
-args = parser.parse_args()
 
-posthog.debug = True
+def load_env_file():
+    """Load environment variables from .env file if it exists."""
+    env_path = os.path.join(os.path.dirname(__file__), ".env")
+    if os.path.exists(env_path):
+        with open(env_path, "r") as f:
+            for line in f:
+                line = line.strip()
+                if line and not line.startswith("#") and "=" in line:
+                    key, value = line.split("=", 1)
+                    os.environ.setdefault(key.strip(), value.strip())
 
-# You can find this key on the /setup page in PostHog
-posthog.project_api_key = "phc_gtWmTq3Pgl06u4sZY3TRcoQfp42yfuXHKoe8ZVSR6Kh"
-posthog.personal_api_key = "phx_fiRCOQkTA3o2ePSdLrFDAILLHjMu2Mv52vUi8MNruIm"
 
-# Where you host PostHog, with no trailing /.
-# You can remove this line if you're using posthog.com
-posthog.host = "http://localhost:8000"
-posthog.poll_interval = 10
+# Load .env file if it exists
+load_env_file()
 
-print(
-    posthog.feature_enabled(
-        args.flag,  # Use the flag from command line arguments
-        "12345",
-        groups={"organization": str("0182ee91-8ef7-0000-4cb9-fedc5f00926a")},
-        group_properties={
-            "organization": {
-                "id": "0182ee91-8ef7-0000-4cb9-fedc5f00926a",
-                "created_at": "2022-06-30 11:44:52.984121+00:00",
-            }
-        },
+# Get configuration
+project_key = os.getenv("POSTHOG_PROJECT_API_KEY", "")
+personal_api_key = os.getenv("POSTHOG_PERSONAL_API_KEY", "")
+host = os.getenv("POSTHOG_HOST", "http://localhost:8000")
+
+# Check if credentials are provided
+if not project_key or not personal_api_key:
+    print("❌ Missing PostHog credentials!")
+    print(
+        "   Please set POSTHOG_PROJECT_API_KEY and POSTHOG_PERSONAL_API_KEY environment variables"
+    )
+    print("   or copy .env.example to .env and fill in your values")
+    exit(1)
+
+# Test authentication before proceeding
+print("🔑 Testing PostHog authentication...")
+
+try:
+    # Configure PostHog with credentials
+    posthog.debug = False  # Keep quiet during auth test
+    posthog.api_key = project_key
+    posthog.project_api_key = project_key
+    posthog.personal_api_key = personal_api_key
+    posthog.host = host
+    posthog.poll_interval = 10
+
+    # Test by attempting to get feature flags (this validates both keys)
+    # This will fail if credentials are invalid
+    test_flags = posthog.get_all_flags("test_user", only_evaluate_locally=True)
+
+    # If we get here without exception, credentials work
+    print("✅ Authentication successful!")
+    print(f"   Project API Key: {project_key[:9]}...")
+    print("   Personal API Key: [REDACTED]")
+    print(f"   Host: {host}\n\n")
+
+except Exception as e:
+    print("❌ Authentication failed!")
+    print(f"   Error: {e}")
+    print("\n   Please check your credentials:")
+    print("   - POSTHOG_PROJECT_API_KEY: Project API key from PostHog settings")
+    print(
+        "   - POSTHOG_PERSONAL_API_KEY: Personal API key (required for local evaluation)"
+    )
+    print("   - POSTHOG_HOST: Your PostHog instance URL")
+    exit(1)
+
+# Display menu and get user choice
+print("🚀 PostHog Python SDK Demo - Choose an example to run:\n")
+print("1. Identify and capture examples")
+print("2. Feature flag local evaluation examples")
+print("3. Feature flag payload examples")
+print("4. Flag dependencies examples")
+print("5. Context management and tagging examples")
+print("6. Run all examples")
+print("7. Exit")
+choice = input("\nEnter your choice (1-7): ").strip()
+
+if choice == "1":
+    print("\n" + "=" * 60)
+    print("IDENTIFY AND CAPTURE EXAMPLES")
+    print("=" * 60)
+
+    posthog.debug = True
+
+    # Capture an event
+    print("📊 Capturing events...")
+    posthog.capture(
+        "event",
+        distinct_id="distinct_id",
+        properties={"property1": "value", "property2": "value"},
+        send_feature_flags=True,
+    )
+
+    # Alias a previous distinct id with a new one
+    print("🔗 Creating alias...")
+    posthog.alias("distinct_id", "new_distinct_id")
+
+    posthog.capture(
+        "event2",
+        distinct_id="new_distinct_id",
+        properties={"property1": "value", "property2": "value"},
+    )
+    posthog.capture(
+        "event-with-groups",
+        distinct_id="new_distinct_id",
+        properties={"property1": "value", "property2": "value"},
+        groups={"company": "id:5"},
+    )
+
+    # Add properties to the person
+    print("👤 Identifying user...")
+    posthog.set(
+        distinct_id="new_distinct_id", properties={"email": "something@something.com"}
+    )
+
+    # Add properties to a group
+    print("🏢 Identifying group...")
+    posthog.group_identify("company", "id:5", {"employees": 11})
+
+    # Properties set only once to the person
+    print("🔒 Setting properties once...")
+    posthog.set_once(
+        distinct_id="new_distinct_id", properties={"self_serve_signup": True}
+    )
+
+    # This will not change the property (because it was already set)
+    posthog.set_once(
+        distinct_id="new_distinct_id", properties={"self_serve_signup": False}
+    )
+
+    print("🔄 Updating properties...")
+    posthog.set(distinct_id="new_distinct_id", properties={"current_browser": "Chrome"})
+    posthog.set(
+        distinct_id="new_distinct_id", properties={"current_browser": "Firefox"}
+    )
+
+elif choice == "2":
+    print("\n" + "=" * 60)
+    print("FEATURE FLAG LOCAL EVALUATION EXAMPLES")
+    print("=" * 60)
+
+    posthog.debug = True
+
+    print("🏁 Testing basic feature flags...")
+    print(
+        f"beta-feature for 'distinct_id': {posthog.feature_enabled('beta-feature', 'distinct_id')}"
+    )
+    print(
+        f"beta-feature for 'new_distinct_id': {posthog.feature_enabled('beta-feature', 'new_distinct_id')}"
+    )
+    print(
+        f"beta-feature with groups: {posthog.feature_enabled('beta-feature-groups', 'distinct_id', groups={'company': 'id:5'})}"
+    )
+
+    print("\n🌍 Testing location-based flags...")
+    # Assume test-flag has `City Name = Sydney` as a person property set
+    print(
+        f"Sydney user: {posthog.feature_enabled('test-flag', 'random_id_12345', person_properties={'$geoip_city_name': 'Sydney'})}"
+    )
+
+    print(
+        f"Sydney user (local only): {posthog.feature_enabled('test-flag', 'distinct_id_random_22', person_properties={'$geoip_city_name': 'Sydney'}, only_evaluate_locally=True)}"
+    )
+
+    print("\n📋 Getting all flags...")
+    print(f"All flags: {posthog.get_all_flags('distinct_id_random_22')}")
+    print(
+        f"All flags (local): {posthog.get_all_flags('distinct_id_random_22', only_evaluate_locally=True)}"
+    )
+    print(
+        f"All flags with properties: {posthog.get_all_flags('distinct_id_random_22', person_properties={'$geoip_city_name': 'Sydney'}, only_evaluate_locally=True)}"
+    )
+
+elif choice == "3":
+    print("\n" + "=" * 60)
+    print("FEATURE FLAG PAYLOAD EXAMPLES")
+    print("=" * 60)
+
+    posthog.debug = True
+
+    print("📦 Testing feature flag payloads...")
+    print(
+        f"beta-feature payload: {posthog.get_feature_flag_payload('beta-feature', 'distinct_id')}"
+    )
+    print(
+        f"All flags and payloads: {posthog.get_all_flags_and_payloads('distinct_id')}"
+    )
+    print(
+        f"Remote config payload: {posthog.get_remote_config_payload('encrypted_payload_flag_key')}"
+    )
+
+    # Get feature flag result with all details (enabled, variant, payload, key, reason)
+    print("\n🔍 Getting detailed flag result...")
+    result = posthog.get_feature_flag_result("beta-feature", "distinct_id")
+    if result:
+        print(f"Flag key: {result.key}")
+        print(f"Flag enabled: {result.enabled}")
+        print(f"Variant: {result.variant}")
+        print(f"Payload: {result.payload}")
+        print(f"Reason: {result.reason}")
+        # get_value() returns the variant if it exists, otherwise the enabled value
+        print(f"Value (variant or enabled): {result.get_value()}")
+
+elif choice == "4":
+    print("\n" + "=" * 60)
+    print("FLAG DEPENDENCIES EXAMPLES")
+    print("=" * 60)
+    print("🔗 Testing flag dependencies with local evaluation...")
+    print(
+        "   Flag structure: 'test-flag-dependency' depends on 'beta-feature' being enabled"
+    )
+    print("")
+    print("📋 Required setup (if 'test-flag-dependency' doesn't exist):")
+    print("   1. Create feature flag 'beta-feature':")
+    print("      - Condition: email contains '@example.com'")
+    print("      - Rollout: 100%")
+    print("   2. Create feature flag 'test-flag-dependency':")
+    print("      - Condition: flag 'beta-feature' is enabled")
+    print("      - Rollout: 100%")
+    print("")
+
+    posthog.debug = True
+
+    # Test @example.com user (should satisfy dependency if flags exist)
+    result1 = posthog.feature_enabled(
+        "test-flag-dependency",
+        "example_user",
+        person_properties={"email": "user@example.com"},
         only_evaluate_locally=True,
     )
-)
+    print(f"✅ @example.com user (test-flag-dependency): {result1}")
 
-
-# Capture an event
-posthog.capture(
-    "event",
-    distinct_id="distinct_id",
-    properties={"property1": "value", "property2": "value"},
-    send_feature_flags=True,
-)
-
-print(posthog.feature_enabled("beta-feature", "distinct_id"))
-print(
-    posthog.feature_enabled(
-        "beta-feature-groups", "distinct_id", groups={"company": "id:5"}
-    )
-)
-
-print(posthog.feature_enabled("beta-feature", "distinct_id"))
-
-# get payload
-print(posthog.get_feature_flag_payload("beta-feature", "distinct_id"))
-print(posthog.get_all_flags_and_payloads("distinct_id"))
-
-# get feature flag result with all details (enabled, variant, payload, key, reason)
-result = posthog.get_feature_flag_result("beta-feature", "distinct_id")
-if result:
-    print(f"Flag key: {result.key}")
-    print(f"Flag enabled: {result.enabled}")
-    print(f"Variant: {result.variant}")
-    print(f"Payload: {result.payload}")
-    print(f"Reason: {result.reason}")
-    # get_value() returns the variant if it exists, otherwise the enabled value
-    print(f"Value (variant or enabled): {result.get_value()}")
-
-# Alias a previous distinct id with a new one
-
-posthog.alias("distinct_id", "new_distinct_id")
-
-posthog.capture(
-    "event2",
-    distinct_id="new_distinct_id",
-    properties={"property1": "value", "property2": "value"},
-)
-posthog.capture(
-    "event-with-groups",
-    distinct_id="new_distinct_id",
-    properties={"property1": "value", "property2": "value"},
-    groups={"company": "id:5"},
-)
-
-# # Add properties to the person
-posthog.set(
-    distinct_id="new_distinct_id", properties={"email": "something@something.com"}
-)
-
-# Add properties to a group
-posthog.group_identify("company", "id:5", {"employees": 11})
-
-# properties set only once to the person
-posthog.set_once(distinct_id="new_distinct_id", properties={"self_serve_signup": True})
-
-
-posthog.set_once(
-    distinct_id="new_distinct_id", properties={"self_serve_signup": False}
-)  # this will not change the property (because it was already set)
-
-posthog.set(distinct_id="new_distinct_id", properties={"current_browser": "Chrome"})
-posthog.set(distinct_id="new_distinct_id", properties={"current_browser": "Firefox"})
-
-
-# #############################################################################
-# Make sure you have a personal API key for the examples below
-
-# Local Evaluation
-
-# If flag has City=Sydney, this call doesn't go to `/decide`
-print(
-    posthog.feature_enabled(
-        "test-flag",
-        "distinct_id_random_22",
-        person_properties={"$geoip_city_name": "Sydney"},
-    )
-)
-
-print(
-    posthog.feature_enabled(
-        "test-flag",
-        "distinct_id_random_22",
-        person_properties={"$geoip_city_name": "Sydney"},
+    # Test non-example.com user (dependency should not be satisfied)
+    result2 = posthog.feature_enabled(
+        "test-flag-dependency",
+        "regular_user",
+        person_properties={"email": "user@other.com"},
         only_evaluate_locally=True,
     )
-)
+    print(f"❌ Regular user (test-flag-dependency): {result2}")
 
-
-print(posthog.get_all_flags("distinct_id_random_22"))
-print(posthog.get_all_flags("distinct_id_random_22", only_evaluate_locally=True))
-print(
-    posthog.get_all_flags(
-        "distinct_id_random_22",
-        person_properties={"$geoip_city_name": "Sydney"},
+    # Test beta-feature directly for comparison
+    beta1 = posthog.feature_enabled(
+        "beta-feature",
+        "example_user",
+        person_properties={"email": "user@example.com"},
         only_evaluate_locally=True,
     )
-)
-print(posthog.get_remote_config_payload("encrypted_payload_flag_key"))
+    beta2 = posthog.feature_enabled(
+        "beta-feature",
+        "regular_user",
+        person_properties={"email": "user@other.com"},
+        only_evaluate_locally=True,
+    )
+    print(f"📊 Beta feature comparison - @example.com: {beta1}, regular: {beta2}")
 
+    print("\n🎯 Results Summary:")
+    print(
+        f"   - Flag dependencies evaluated locally: {'✅ YES' if result1 != result2 else '❌ NO'}"
+    )
+    print("   - Zero API calls needed: ✅ YES (all evaluated locally)")
+    print("   - Python SDK supports flag dependencies: ✅ YES")
 
-# You can add tags to a context, and these are automatically added to any events (including exceptions) captured
-# within that context.
+    print("\n" + "-" * 60)
+    print("PRODUCTION-STYLE MULTIVARIATE DEPENDENCY CHAIN")
+    print("-" * 60)
+    print("🔗 Testing complex multivariate flag dependencies...")
+    print(
+        "   Structure: multivariate-root-flag -> multivariate-intermediate-flag -> multivariate-leaf-flag"
+    )
+    print("")
+    print("📋 Required setup (if flags don't exist):")
+    print(
+        "   1. Create 'multivariate-leaf-flag' with fruit variants (pineapple, mango, papaya, kiwi)"
+    )
+    print("      - pineapple: email = 'pineapple@example.com'")
+    print("      - mango: email = 'mango@example.com'")
+    print(
+        "   2. Create 'multivariate-intermediate-flag' with color variants (blue, red)"
+    )
+    print("      - blue: depends on multivariate-leaf-flag = 'pineapple'")
+    print("      - red: depends on multivariate-leaf-flag = 'mango'")
+    print(
+        "   3. Create 'multivariate-root-flag' with show variants (breaking-bad, the-wire)"
+    )
+    print("      - breaking-bad: depends on multivariate-intermediate-flag = 'blue'")
+    print("      - the-wire: depends on multivariate-intermediate-flag = 'red'")
+    print("")
 
-# You can enter a new context using a with statement. Any exceptions thrown in the context will be captured,
-# and tagged with the context tags. Other events captured will also be tagged with the context tags. By default,
-# the new context inherits tags from the parent context.
-with posthog.new_context():
-    posthog.tag("transaction_id", "abc123")
-    posthog.tag("some_arbitrary_value", {"tags": "can be dicts"})
+    # Test pineapple -> blue -> breaking-bad chain
+    dependent_result3 = posthog.get_feature_flag(
+        "multivariate-root-flag",
+        "regular_user",
+        person_properties={"email": "pineapple@example.com"},
+        only_evaluate_locally=True,
+    )
+    if str(dependent_result3) != "breaking-bad":
+        print(
+            f"     ❌ Something went wrong evaluating 'multivariate-root-flag' with pineapple@example.com. Expected 'breaking-bad', got '{dependent_result3}'"
+        )
+    else:
+        print("✅ 'multivariate-root-flag' with email pineapple@example.com succeeded")
 
-    # This event will be captured with the tags set above
-    posthog.capture("order_processed")
-    # This exception will be captured with the tags set above
-    raise Exception("Order processing failed")
+    # Test mango -> red -> the-wire chain
+    dependent_result4 = posthog.get_feature_flag(
+        "multivariate-root-flag",
+        "regular_user",
+        person_properties={"email": "mango@example.com"},
+        only_evaluate_locally=True,
+    )
+    if str(dependent_result4) != "the-wire":
+        print(
+            f"     ❌ Something went wrong evaluating multivariate-root-flag with mango@example.com. Expected 'the-wire', got '{dependent_result4}'"
+        )
+    else:
+        print("✅ 'multivariate-root-flag' with email mango@example.com succeeded")
 
+    # Show the complete chain evaluation
+    print("\n🔍 Complete dependency chain evaluation:")
+    for email, expected_chain in [
+        ("pineapple@example.com", ["pineapple", "blue", "breaking-bad"]),
+        ("mango@example.com", ["mango", "red", "the-wire"]),
+    ]:
+        leaf = posthog.get_feature_flag(
+            "multivariate-leaf-flag",
+            "regular_user",
+            person_properties={"email": email},
+            only_evaluate_locally=True,
+        )
+        intermediate = posthog.get_feature_flag(
+            "multivariate-intermediate-flag",
+            "regular_user",
+            person_properties={"email": email},
+            only_evaluate_locally=True,
+        )
+        root = posthog.get_feature_flag(
+            "multivariate-root-flag",
+            "regular_user",
+            person_properties={"email": email},
+            only_evaluate_locally=True,
+        )
 
-# Use fresh=True to start with a clean context (no inherited tags)
-with posthog.new_context(fresh=True):
-    posthog.tag("session_id", "xyz789")
-    # Only session_id tag will be present, no inherited tags
-    raise Exception("Session handling failed")
+        actual_chain = [str(leaf), str(intermediate), str(root)]
+        chain_success = actual_chain == expected_chain
 
+        print(f"   📧 {email}:")
+        print(f"      Expected: {' -> '.join(map(str, expected_chain))}")
+        print(f"      Actual:   {' -> '.join(map(str, actual_chain))}")
+        print(f"      Status:   {'✅ SUCCESS' if chain_success else '❌ FAILED'}")
 
-# You can also use the `@posthog.scoped()` decorator to enter a new context.
-# By default, it inherits tags from the parent context
-@posthog.scoped()
-def process_order(order_id):
-    posthog.tag("order_id", order_id)
-    # Exception will be captured and tagged automatically
-    raise Exception("Order processing failed")
+    print("\n🎯 Multivariate Chain Summary:")
+    print("   - Complex dependency chains: ✅ SUPPORTED")
+    print("   - Multivariate flag dependencies: ✅ SUPPORTED")
+    print("   - Local evaluation of chains: ✅ WORKING")
 
+elif choice == "5":
+    print("\n" + "=" * 60)
+    print("CONTEXT MANAGEMENT AND TAGGING EXAMPLES")
+    print("=" * 60)
 
-# Use fresh=True to start with a clean context (no inherited tags)
-@posthog.scoped(fresh=True)
-def process_payment(payment_id):
-    posthog.tag("payment_id", payment_id)
-    # Only payment_id tag will be present, no inherited tags
-    raise Exception("Payment processing failed")
+    posthog.debug = True
 
+    print("🏷️ Testing context management...")
+    print(
+        "You can add tags to a context, and these are automatically added to any events captured within that context."
+    )
+
+    # You can enter a new context using a with statement. Any exceptions thrown in the context will be captured,
+    # and tagged with the context tags. Other events captured will also be tagged with the context tags. By default,
+    # the new context inherits tags from the parent context.
+    try:
+        with posthog.new_context():
+            posthog.tag("transaction_id", "abc123")
+            posthog.tag("some_arbitrary_value", {"tags": "can be dicts"})
+
+            # This event will be captured with the tags set above
+            posthog.capture("order_processed")
+            print("✅ Event captured with inherited context tags")
+            # This exception will be captured with the tags set above
+            # raise Exception("Order processing failed")
+    except Exception as e:
+        print(f"Exception captured: {e}")
+
+    # Use fresh=True to start with a clean context (no inherited tags)
+    try:
+        with posthog.new_context(fresh=True):
+            posthog.tag("session_id", "xyz789")
+            # Only session_id tag will be present, no inherited tags
+            posthog.capture("session_event")
+            print("✅ Event captured with fresh context tags")
+            # raise Exception("Session handling failed")
+    except Exception as e:
+        print(f"Exception captured: {e}")
+
+    # You can also use the `@posthog.scoped()` decorator to enter a new context.
+    # By default, it inherits tags from the parent context
+    @posthog.scoped()
+    def process_order(order_id):
+        posthog.tag("order_id", order_id)
+        posthog.capture("order_step_completed")
+        print(f"✅ Order {order_id} processed with scoped context")
+        # Exception will be captured and tagged automatically
+        # raise Exception("Order processing failed")
+
+    # Use fresh=True to start with a clean context (no inherited tags)
+    @posthog.scoped(fresh=True)
+    def process_payment(payment_id):
+        posthog.tag("payment_id", payment_id)
+        posthog.capture("payment_processed")
+        print(f"✅ Payment {payment_id} processed with fresh scoped context")
+        # Only payment_id tag will be present, no inherited tags
+        # raise Exception("Payment processing failed")
+
+    process_order("12345")
+    process_payment("67890")
+
+elif choice == "6":
+    print("\n🔄 Running all examples...")
+
+    # Run example 1
+    print(f"\n{'🔸' * 20} IDENTIFY AND CAPTURE {'🔸' * 20}")
+    posthog.debug = True
+    print("📊 Capturing events...")
+    posthog.capture(
+        "event",
+        distinct_id="distinct_id",
+        properties={"property1": "value", "property2": "value"},
+        send_feature_flags=True,
+    )
+    print("🔗 Creating alias...")
+    posthog.alias("distinct_id", "new_distinct_id")
+    print("👤 Identifying user...")
+    posthog.set(
+        distinct_id="new_distinct_id", properties={"email": "something@something.com"}
+    )
+
+    # Run example 2
+    print(f"\n{'🔸' * 20} FEATURE FLAGS {'🔸' * 20}")
+    print("🏁 Testing basic feature flags...")
+    print(f"beta-feature: {posthog.feature_enabled('beta-feature', 'distinct_id')}")
+    print(
+        f"Sydney user: {posthog.feature_enabled('test-flag', 'random_id_12345', person_properties={'$geoip_city_name': 'Sydney'})}"
+    )
+
+    # Run example 3
+    print(f"\n{'🔸' * 20} PAYLOADS {'🔸' * 20}")
+    print("📦 Testing payloads...")
+    print(f"Payload: {posthog.get_feature_flag_payload('beta-feature', 'distinct_id')}")
+
+    # Run example 4
+    print(f"\n{'🔸' * 20} FLAG DEPENDENCIES {'🔸' * 20}")
+    print("🔗 Testing flag dependencies...")
+    result1 = posthog.feature_enabled(
+        "test-flag-dependency",
+        "demo_user",
+        person_properties={"email": "user@example.com"},
+        only_evaluate_locally=True,
+    )
+    result2 = posthog.feature_enabled(
+        "test-flag-dependency",
+        "demo_user2",
+        person_properties={"email": "user@other.com"},
+        only_evaluate_locally=True,
+    )
+    print(f"✅ @example.com user: {result1}, regular user: {result2}")
+
+    # Run example 5
+    print(f"\n{'🔸' * 20} CONTEXT MANAGEMENT {'🔸' * 20}")
+    print("🏷️ Testing context management...")
+    with posthog.new_context():
+        posthog.tag("demo_run", "all_examples")
+        posthog.capture("demo_completed")
+        print("✅ Demo completed with context tags")
+
+elif choice == "7":
+    print("👋 Goodbye!")
+    posthog.shutdown()
+    exit()
+
+else:
+    print("❌ Invalid choice. Please run again and select 1-7.")
+    posthog.shutdown()
+    exit()
+
+print("\n" + "=" * 60)
+print("✅ Example completed!")
+print("=" * 60)
 
 posthog.shutdown()

mypy-baseline.txt

@@ -36,11 +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]
-example.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/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, cast
+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
 
 
@@ -42,6 +51,12 @@ class Client:
     def __init__(
         self,
         api_key: Optional[str] = None,
+        vertexai: Optional[bool] = None,
+        credentials: Optional[Any] = None,
+        project: Optional[str] = None,
+        location: Optional[str] = None,
+        debug_config: Optional[Any] = None,
+        http_options: Optional[Any] = None,
         posthog_client: Optional[PostHogClient] = None,
         posthog_distinct_id: Optional[str] = None,
         posthog_properties: Optional[Dict[str, Any]] = None,
@@ -51,7 +66,13 @@ class Client:
     ):
         """
         Args:
-            api_key: Google AI API key. If not provided, will use GOOGLE_API_KEY or API_KEY environment variable
+            api_key: Google AI API key. If not provided, will use GOOGLE_API_KEY or API_KEY environment variable (not required for Vertex AI)
+            vertexai: Whether to use Vertex AI authentication
+            credentials: Vertex AI credentials object
+            project: GCP project ID for Vertex AI
+            location: GCP location for Vertex AI
+            debug_config: Debug configuration for the client
+            http_options: HTTP options for the client
             posthog_client: PostHog client for tracking usage
             posthog_distinct_id: Default distinct ID for all calls (can be overridden per call)
             posthog_properties: Default properties for all calls (can be overridden per call)
@@ -59,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:
@@ -66,6 +88,12 @@ class Client:
 
         self.models = Models(
             api_key=api_key,
+            vertexai=vertexai,
+            credentials=credentials,
+            project=project,
+            location=location,
+            debug_config=debug_config,
+            http_options=http_options,
             posthog_client=self._ph_client,
             posthog_distinct_id=posthog_distinct_id,
             posthog_properties=posthog_properties,
@@ -85,6 +113,12 @@ class Models:
     def __init__(
         self,
         api_key: Optional[str] = None,
+        vertexai: Optional[bool] = None,
+        credentials: Optional[Any] = None,
+        project: Optional[str] = None,
+        location: Optional[str] = None,
+        debug_config: Optional[Any] = None,
+        http_options: Optional[Any] = None,
         posthog_client: Optional[PostHogClient] = None,
         posthog_distinct_id: Optional[str] = None,
         posthog_properties: Optional[Dict[str, Any]] = None,
@@ -94,7 +128,13 @@ class Models:
     ):
         """
         Args:
-            api_key: Google AI API key. If not provided, will use GOOGLE_API_KEY or API_KEY environment variable
+            api_key: Google AI API key. If not provided, will use GOOGLE_API_KEY or API_KEY environment variable (not required for Vertex AI)
+            vertexai: Whether to use Vertex AI authentication
+            credentials: Vertex AI credentials object
+            project: GCP project ID for Vertex AI
+            location: GCP location for Vertex AI
+            debug_config: Debug configuration for the client
+            http_options: HTTP options for the client
             posthog_client: PostHog client for tracking usage
             posthog_distinct_id: Default distinct ID for all calls
             posthog_properties: Default properties for all calls
@@ -102,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:
@@ -113,16 +154,46 @@ class Models:
         self._default_privacy_mode = posthog_privacy_mode
         self._default_groups = posthog_groups
 
-        # Handle API key - try parameter first, then environment variables
-        if api_key is None:
-            api_key = os.environ.get("GOOGLE_API_KEY") or os.environ.get("API_KEY")
+        # Build genai.Client arguments
+        client_args: Dict[str, Any] = {}
 
-        if api_key is None:
-            raise ValueError(
-                "API key must be provided either as parameter or via GOOGLE_API_KEY/API_KEY environment variable"
-            )
+        # Add Vertex AI parameters if provided
+        if vertexai is not None:
+            client_args["vertexai"] = vertexai
 
-        self._client = genai.Client(api_key=api_key)
+        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
+
+        # Handle API key authentication
+        if vertexai:
+            # For Vertex AI, api_key is optional
+            if api_key is not None:
+                client_args["api_key"] = api_key
+        else:
+            # For non-Vertex AI mode, api_key is required (backwards compatibility)
+            if api_key is None:
+                api_key = os.environ.get("GOOGLE_API_KEY") or os.environ.get("API_KEY")
+
+            if api_key is None:
+                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)
         self._base_url = "https://generativelanguage.googleapis.com"
 
     def _merge_posthog_params(
@@ -134,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
@@ -149,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)
 
@@ -184,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(
@@ -222,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}
@@ -233,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,
@@ -264,7 +337,7 @@ class Models:
                     kwargs,
                     usage_stats,
                     latency,
-                    output,
+                    accumulated_content,
                 )
 
         return generator()
@@ -279,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,7 +481,7 @@ 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,
@@ -550,7 +551,7 @@ 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,

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

@@ -1,6 +1,8 @@
 import time
 import uuid
-from typing import Any, Dict, List, Optional, cast
+from typing import Any, Dict, List, Optional
+
+from posthog.ai.types import TokenUsage
 
 try:
     import openai
@@ -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/client.py

@@ -3,7 +3,7 @@ import logging
 import os
 import sys
 from datetime import datetime, timedelta
-from typing import Any, Dict, Optional, Union
+from typing import Any, Callable, Dict, Optional, Union
 from typing_extensions import Unpack
 from uuid import uuid4
 
@@ -44,6 +44,7 @@ from posthog.types import (
     FlagsAndPayloads,
     FlagsResponse,
     FlagValue,
+    SendFeatureFlagsOptions,
     normalize_flags_response,
     to_flags_and_payloads,
     to_payloads,
@@ -87,6 +88,7 @@ def add_context_tags(properties):
     current_context = _get_current_context()
     if current_context:
         context_tags = current_context.collect_tags()
+        properties["$context_tags"] = set(context_tags.keys())
         # We want explicitly passed properties to override context tags
         context_tags.update(properties)
         properties = context_tags
@@ -97,6 +99,34 @@ def add_context_tags(properties):
     return properties
 
 
+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.
@@ -312,6 +342,7 @@ class Client(object):
         person_properties=None,
         group_properties=None,
         disable_geoip=None,
+        flag_keys_to_evaluate: Optional[list[str]] = None,
     ) -> dict[str, Union[bool, str]]:
         """
         Get feature flag variants for a user by calling decide.
@@ -322,12 +353,19 @@ class Client(object):
             person_properties: A dictionary of person properties.
             group_properties: A dictionary of group properties.
             disable_geoip: Whether to disable GeoIP for this request.
+            flag_keys_to_evaluate: A list of specific flag keys to evaluate. If provided,
+                only these flags will be evaluated, improving performance.
 
         Category:
-            Feature Flags
+            Feature flags
         """
         resp_data = self.get_flags_decision(
-            distinct_id, groups, person_properties, group_properties, disable_geoip
+            distinct_id,
+            groups,
+            person_properties,
+            group_properties,
+            disable_geoip,
+            flag_keys_to_evaluate,
         )
         return to_values(resp_data) or {}
 
@@ -338,6 +376,7 @@ class Client(object):
         person_properties=None,
         group_properties=None,
         disable_geoip=None,
+        flag_keys_to_evaluate: Optional[list[str]] = None,
     ) -> dict[str, str]:
         """
         Get feature flag payloads for a user by calling decide.
@@ -348,6 +387,8 @@ class Client(object):
             person_properties: A dictionary of person properties.
             group_properties: A dictionary of group properties.
             disable_geoip: Whether to disable GeoIP for this request.
+            flag_keys_to_evaluate: A list of specific flag keys to evaluate. If provided,
+                only these flags will be evaluated, improving performance.
 
         Examples:
             ```python
@@ -355,10 +396,15 @@ class Client(object):
             ```
 
         Category:
-            Feature Flags
+            Feature flags
         """
         resp_data = self.get_flags_decision(
-            distinct_id, groups, person_properties, group_properties, disable_geoip
+            distinct_id,
+            groups,
+            person_properties,
+            group_properties,
+            disable_geoip,
+            flag_keys_to_evaluate,
         )
         return to_payloads(resp_data) or {}
 
@@ -369,6 +415,7 @@ class Client(object):
         person_properties=None,
         group_properties=None,
         disable_geoip=None,
+        flag_keys_to_evaluate: Optional[list[str]] = None,
     ) -> FlagsAndPayloads:
         """
         Get feature flags and payloads for a user by calling decide.
@@ -379,6 +426,8 @@ class Client(object):
             person_properties: A dictionary of person properties.
             group_properties: A dictionary of group properties.
             disable_geoip: Whether to disable GeoIP for this request.
+            flag_keys_to_evaluate: A list of specific flag keys to evaluate. If provided,
+                only these flags will be evaluated, improving performance.
 
         Examples:
             ```python
@@ -386,10 +435,15 @@ class Client(object):
             ```
 
         Category:
-            Feature Flags
+            Feature flags
         """
         resp = self.get_flags_decision(
-            distinct_id, groups, person_properties, group_properties, disable_geoip
+            distinct_id,
+            groups,
+            person_properties,
+            group_properties,
+            disable_geoip,
+            flag_keys_to_evaluate,
         )
         return to_flags_and_payloads(resp)
 
@@ -400,6 +454,7 @@ class Client(object):
         person_properties=None,
         group_properties=None,
         disable_geoip=None,
+        flag_keys_to_evaluate: Optional[list[str]] = None,
     ) -> FlagsResponse:
         """
         Get feature flags decision.
@@ -410,6 +465,8 @@ class Client(object):
             person_properties: A dictionary of person properties.
             group_properties: A dictionary of group properties.
             disable_geoip: Whether to disable GeoIP for this request.
+            flag_keys_to_evaluate: A list of specific flag keys to evaluate. If provided,
+                only these flags will be evaluated, improving performance.
 
         Examples:
             ```python
@@ -417,7 +474,7 @@ class Client(object):
             ```
 
         Category:
-            Feature Flags
+            Feature flags
         """
         groups = groups or {}
         person_properties = person_properties or {}
@@ -440,6 +497,9 @@ class Client(object):
             "geoip_disable": disable_geoip,
         }
 
+        if flag_keys_to_evaluate:
+            request_data["flag_keys_to_evaluate"] = flag_keys_to_evaluate
+
         resp_data = flags(
             self.api_key,
             self.host,
@@ -449,6 +509,7 @@ class Client(object):
 
         return normalize_flags_response(resp_data)
 
+    @no_throw()
     def capture(
         self, event: str, **kwargs: Unpack[OptionalCaptureArgs]
     ) -> Optional[str]:
@@ -544,6 +605,7 @@ class Client(object):
                         group_properties=flag_options["group_properties"],
                         disable_geoip=disable_geoip,
                         only_evaluate_locally=True,
+                        flag_keys_to_evaluate=flag_options["flag_keys_filter"],
                     )
                 else:
                     # Default behavior - use remote evaluation
@@ -553,6 +615,7 @@ class Client(object):
                         person_properties=flag_options["person_properties"],
                         group_properties=flag_options["group_properties"],
                         disable_geoip=disable_geoip,
+                        flag_keys_to_evaluate=flag_options["flag_keys_filter"],
                     )
             except Exception as e:
                 self.log.exception(
@@ -585,7 +648,7 @@ class Client(object):
 
         return self._enqueue(msg, disable_geoip)
 
-    def _parse_send_feature_flags(self, send_feature_flags) -> dict:
+    def _parse_send_feature_flags(self, send_feature_flags) -> SendFeatureFlagsOptions:
         """
         Parse and normalize send_feature_flags parameter into a standard format.
 
@@ -593,8 +656,8 @@ class Client(object):
             send_feature_flags: Either bool or SendFeatureFlagsOptions dict
 
         Returns:
-            dict: Normalized options with keys: should_send, only_evaluate_locally,
-                  person_properties, group_properties
+            SendFeatureFlagsOptions: Normalized options with keys: should_send, only_evaluate_locally,
+                  person_properties, group_properties, flag_keys_filter
 
         Raises:
             TypeError: If send_feature_flags is not bool or dict
@@ -607,6 +670,7 @@ class Client(object):
                 ),
                 "person_properties": send_feature_flags.get("person_properties"),
                 "group_properties": send_feature_flags.get("group_properties"),
+                "flag_keys_filter": send_feature_flags.get("flag_keys_filter"),
             }
         elif isinstance(send_feature_flags, bool):
             return {
@@ -614,6 +678,7 @@ class Client(object):
                 "only_evaluate_locally": None,
                 "person_properties": None,
                 "group_properties": None,
+                "flag_keys_filter": None,
             }
         else:
             raise TypeError(
@@ -621,6 +686,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.
@@ -654,6 +720,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)
@@ -680,6 +748,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.
@@ -698,6 +767,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)
@@ -723,6 +794,7 @@ class Client(object):
 
         return self._enqueue(msg, disable_geoip)
 
+    @no_throw()
     def group_identify(
         self,
         group_type: str,
@@ -755,6 +827,8 @@ class Client(object):
 
         Category:
             Identification
+
+        Note: This method will not raise exceptions. Errors are logged.
         """
         properties = properties or {}
 
@@ -779,6 +853,7 @@ class Client(object):
 
         return self._enqueue(msg, disable_geoip)
 
+    @no_throw()
     def alias(
         self,
         previous_id: str,
@@ -804,6 +879,8 @@ class Client(object):
 
         Category:
             Identification
+
+        Note: This method will not raise exceptions. Errors are logged.
         """
         (distinct_id, personless) = get_identity_state(distinct_id)
 
@@ -896,7 +973,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,
             }
 
@@ -1133,7 +1209,7 @@ class Client(object):
             ```
 
         Category:
-            Feature Flags
+            Feature flags
         """
         if not self.personal_api_key:
             self.log.warning(
@@ -1168,6 +1244,9 @@ class Client(object):
         person_properties = person_properties or {}
         group_properties = group_properties or {}
 
+        # Create evaluation cache for flag dependencies
+        evaluation_cache: dict[str, Optional[FlagValue]] = {}
+
         if feature_flag.get("ensure_experience_continuity", False):
             raise InconclusiveMatchError("Flag has experience continuity enabled")
 
@@ -1183,12 +1262,12 @@ class Client(object):
                 self.log.warning(
                     f"[FEATURE FLAGS] Unknown group type index {aggregation_group_type_index} for feature flag {feature_flag['key']}"
                 )
-                # failover to `/decide/`
+                # failover to `/flags`
                 raise InconclusiveMatchError("Flag has unknown group type index")
 
             if group_name not in groups:
                 # Group flags are never enabled in `groups` aren't passed in
-                # don't failover to `/decide/`, since response will be the same
+                # don't failover to `/flags`, since response will be the same
                 if warn_on_unknown_groups:
                     self.log.warning(
                         f"[FEATURE FLAGS] Can't compute group feature flag: {feature_flag['key']} without group names passed in"
@@ -1201,11 +1280,20 @@ class Client(object):
 
             focused_group_properties = group_properties[group_name]
             return match_feature_flag_properties(
-                feature_flag, groups[group_name], focused_group_properties
+                feature_flag,
+                groups[group_name],
+                focused_group_properties,
+                self.feature_flags_by_key,
+                evaluation_cache,
             )
         else:
             return match_feature_flag_properties(
-                feature_flag, distinct_id, person_properties, self.cohorts
+                feature_flag,
+                distinct_id,
+                person_properties,
+                self.cohorts,
+                self.feature_flags_by_key,
+                evaluation_cache,
             )
 
     def feature_enabled(
@@ -1243,7 +1331,7 @@ class Client(object):
             ```
 
         Category:
-            Feature Flags
+            Feature flags
         """
         response = self.get_feature_flag(
             key,
@@ -1316,7 +1404,7 @@ class Client(object):
                 )
         elif not only_evaluate_locally:
             try:
-                flag_details, request_id = self._get_feature_flag_details_from_decide(
+                flag_details, request_id = self._get_feature_flag_details_from_server(
                     key,
                     distinct_id,
                     groups,
@@ -1451,7 +1539,7 @@ class Client(object):
             ```
 
         Category:
-            Feature Flags
+            Feature flags
         """
         feature_flag_result = self.get_feature_flag_result(
             key,
@@ -1541,7 +1629,7 @@ class Client(object):
             ```
 
         Category:
-            Feature Flags
+            Feature flags
         """
         feature_flag_result = self._get_feature_flag_result(
             key,
@@ -1556,7 +1644,7 @@ class Client(object):
         )
         return feature_flag_result.payload if feature_flag_result else None
 
-    def _get_feature_flag_details_from_decide(
+    def _get_feature_flag_details_from_server(
         self,
         key: str,
         distinct_id: ID_TYPES,
@@ -1566,10 +1654,15 @@ class Client(object):
         disable_geoip: Optional[bool],
     ) -> tuple[Optional[FeatureFlag], Optional[str]]:
         """
-        Calls /decide and returns the flag details and request id
+        Calls /flags and returns the flag details and request id
         """
         resp_data = self.get_flags_decision(
-            distinct_id, groups, person_properties, group_properties, disable_geoip
+            distinct_id,
+            groups,
+            person_properties,
+            group_properties,
+            disable_geoip,
+            flag_keys_to_evaluate=[key],
         )
         request_id = resp_data.get("requestId")
         flags = resp_data.get("flags")
@@ -1644,6 +1737,7 @@ class Client(object):
         try:
             return remote_config(
                 self.personal_api_key,
+                self.api_key,
                 self.host,
                 key,
                 timeout=self.feature_flags_request_timeout_seconds,
@@ -1684,6 +1778,7 @@ class Client(object):
         group_properties=None,
         only_evaluate_locally=False,
         disable_geoip=None,
+        flag_keys_to_evaluate: Optional[list[str]] = None,
     ) -> Optional[dict[str, Union[bool, str]]]:
         """
         Get all feature flags for a user.
@@ -1695,6 +1790,8 @@ class Client(object):
             group_properties: A dictionary of group properties.
             only_evaluate_locally: Whether to only evaluate locally.
             disable_geoip: Whether to disable GeoIP for this request.
+            flag_keys_to_evaluate: A list of specific flag keys to evaluate. If provided,
+                only these flags will be evaluated, improving performance.
 
         Examples:
             ```python
@@ -1702,7 +1799,7 @@ class Client(object):
             ```
 
         Category:
-            Feature Flags
+            Feature flags
         """
         response = self.get_all_flags_and_payloads(
             distinct_id,
@@ -1711,6 +1808,7 @@ class Client(object):
             group_properties=group_properties,
             only_evaluate_locally=only_evaluate_locally,
             disable_geoip=disable_geoip,
+            flag_keys_to_evaluate=flag_keys_to_evaluate,
         )
 
         return response["featureFlags"]
@@ -1724,6 +1822,7 @@ class Client(object):
         group_properties=None,
         only_evaluate_locally=False,
         disable_geoip=None,
+        flag_keys_to_evaluate: Optional[list[str]] = None,
     ) -> FlagsAndPayloads:
         """
         Get all feature flags and their payloads for a user.
@@ -1735,6 +1834,8 @@ class Client(object):
             group_properties: A dictionary of group properties.
             only_evaluate_locally: Whether to only evaluate locally.
             disable_geoip: Whether to disable GeoIP for this request.
+            flag_keys_to_evaluate: A list of specific flag keys to evaluate. If provided,
+                only these flags will be evaluated, improving performance.
 
         Examples:
             ```python
@@ -1742,7 +1843,7 @@ class Client(object):
             ```
 
         Category:
-            Feature Flags
+            Feature flags
         """
         if self.disabled:
             return {"featureFlags": None, "featureFlagPayloads": None}
@@ -1753,14 +1854,15 @@ 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,
             group_properties=group_properties,
+            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,
@@ -1768,6 +1870,7 @@ class Client(object):
                     person_properties=person_properties,
                     group_properties=group_properties,
                     disable_geoip=disable_geoip,
+                    flag_keys_to_evaluate=flag_keys_to_evaluate,
                 )
                 return to_flags_and_payloads(decide_response)
             except Exception as e:
@@ -1785,6 +1888,7 @@ class Client(object):
         person_properties=None,
         group_properties=None,
         warn_on_unknown_groups=False,
+        flag_keys_to_evaluate: Optional[list[str]] = None,
     ) -> tuple[FlagsAndPayloads, bool]:
         person_properties = person_properties or {}
         group_properties = group_properties or {}
@@ -1794,10 +1898,18 @@ 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:
-            for flag in self.feature_flags:
+            # Filter flags based on flag_keys_to_evaluate if provided
+            flags_to_process = self.feature_flags
+            if flag_keys_to_evaluate:
+                flag_keys_set = set(flag_keys_to_evaluate)
+                flags_to_process = [
+                    flag for flag in self.feature_flags if flag["key"] in flag_keys_set
+                ]
+
+            for flag in flags_to_process:
                 try:
                     flags[flag["key"]] = self._compute_flag_locally(
                         flag,
@@ -1813,20 +1925,20 @@ class Client(object):
                     if matched_payload is not None:
                         payloads[flag["key"]] = matched_payload
                 except InconclusiveMatchError:
-                    # No need to log this, since it's just telling us to fall back to `/decide`
-                    fallback_to_decide = True
+                    # No need to log this, since it's just telling us to fall back to `/flags`
+                    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.
@@ -1937,7 +2049,7 @@ class Client(object):
             for group_name in groups:
                 all_group_properties[group_name] = {
                     "$group_key": groups[group_name],
-                    **(group_properties.get(group_name) or {}),
+                    **((group_properties or {}).get(group_name) or {}),
                 }
 
         return all_person_properties, all_group_properties

posthog/feature_flags.py

@@ -55,8 +55,161 @@ def variant_lookup_table(feature_flag):
     return lookup_table
 
 
+def evaluate_flag_dependency(
+    property, flags_by_key, evaluation_cache, distinct_id, properties, cohort_properties
+):
+    """
+    Evaluate a flag dependency property according to the dependency chain algorithm.
+
+    Args:
+        property: Flag property with type="flag" and dependency_chain
+        flags_by_key: Dictionary of all flags by their key
+        evaluation_cache: Cache for storing evaluation results
+        distinct_id: The distinct ID being evaluated
+        properties: Person properties for evaluation
+        cohort_properties: Cohort properties for evaluation
+
+    Returns:
+        bool: True if all dependencies in the chain evaluate to True, False otherwise
+    """
+    if flags_by_key is None or evaluation_cache is None:
+        # Cannot evaluate flag dependencies without required context
+        raise InconclusiveMatchError(
+            f"Cannot evaluate flag dependency on '{property.get('key', 'unknown')}' without flags_by_key and evaluation_cache"
+        )
+
+    # Check if dependency_chain is present - it should always be provided for flag dependencies
+    if "dependency_chain" not in property:
+        # Missing dependency_chain indicates malformed server data
+        raise InconclusiveMatchError(
+            f"Flag dependency property for '{property.get('key', 'unknown')}' is missing required 'dependency_chain' field"
+        )
+
+    dependency_chain = property["dependency_chain"]
+
+    # Handle circular dependency (empty chain means circular)
+    if len(dependency_chain) == 0:
+        log.debug(f"Circular dependency detected for flag: {property.get('key')}")
+        raise InconclusiveMatchError(
+            f"Circular dependency detected for flag '{property.get('key', 'unknown')}'"
+        )
+
+    # Evaluate all dependencies in the chain order
+    for dep_flag_key in dependency_chain:
+        if dep_flag_key not in evaluation_cache:
+            # Need to evaluate this dependency first
+            dep_flag = flags_by_key.get(dep_flag_key)
+            if not dep_flag:
+                # Missing flag dependency - cannot evaluate locally
+                evaluation_cache[dep_flag_key] = None
+                raise InconclusiveMatchError(
+                    f"Cannot evaluate flag dependency '{dep_flag_key}' - flag not found in local flags"
+                )
+            else:
+                # Check if the flag is active (same check as in client._compute_flag_locally)
+                if not dep_flag.get("active"):
+                    evaluation_cache[dep_flag_key] = False
+                else:
+                    # Recursively evaluate the dependency
+                    try:
+                        dep_result = match_feature_flag_properties(
+                            dep_flag,
+                            distinct_id,
+                            properties,
+                            cohort_properties,
+                            flags_by_key,
+                            evaluation_cache,
+                        )
+                        evaluation_cache[dep_flag_key] = dep_result
+                    except InconclusiveMatchError as e:
+                        # If we can't evaluate a dependency, store None and propagate the error
+                        evaluation_cache[dep_flag_key] = None
+                        raise InconclusiveMatchError(
+                            f"Cannot evaluate flag dependency '{dep_flag_key}': {e}"
+                        ) from e
+
+        # Check the cached result
+        cached_result = evaluation_cache[dep_flag_key]
+        if cached_result is None:
+            # Previously inconclusive - raise error again
+            raise InconclusiveMatchError(
+                f"Flag dependency '{dep_flag_key}' was previously inconclusive"
+            )
+        elif not cached_result:
+            # Definitive False result - dependency failed
+            return False
+
+    # All dependencies in the chain have been evaluated successfully
+    # Now check if the final flag value matches the expected value in the property
+    flag_key = property.get("key")
+    expected_value = property.get("value")
+    operator = property.get("operator", "exact")
+
+    if flag_key and expected_value is not None:
+        # Get the actual value of the flag we're checking
+        actual_value = evaluation_cache.get(flag_key)
+
+        if actual_value is None:
+            # Flag wasn't evaluated - this shouldn't happen if dependency chain is correct
+            raise InconclusiveMatchError(
+                f"Flag '{flag_key}' was not evaluated despite being in dependency chain"
+            )
+
+        # For flag dependencies, we need to compare the actual flag result with expected value
+        # using the flag_evaluates_to operator logic
+        if operator == "flag_evaluates_to":
+            return matches_dependency_value(expected_value, actual_value)
+        else:
+            # This should never happen, but just to be defensive.
+            raise InconclusiveMatchError(
+                f"Flag dependency property for '{property.get('key', 'unknown')}' has invalid operator '{operator}'"
+            )
+
+    # If no value check needed, return True (all dependencies passed)
+    return True
+
+
+def matches_dependency_value(expected_value, actual_value):
+    """
+    Check if the actual flag value matches the expected dependency value.
+
+    This follows the same logic as the C# MatchesDependencyValue function:
+    - String variant case: check for exact match or boolean true
+    - Boolean case: must match expected boolean value
+
+    Args:
+        expected_value: The expected value from the property
+        actual_value: The actual value returned by the flag evaluation
+
+    Returns:
+        bool: True if the values match according to flag dependency rules
+    """
+    # String variant case - check for exact match or boolean true
+    if isinstance(actual_value, str) and len(actual_value) > 0:
+        if isinstance(expected_value, bool):
+            # Any variant matches boolean true
+            return expected_value
+        elif isinstance(expected_value, str):
+            # variants are case-sensitive, hence our comparison is too
+            return actual_value == expected_value
+        else:
+            return False
+
+    # Boolean case - must match expected boolean value
+    elif isinstance(actual_value, bool) and isinstance(expected_value, bool):
+        return actual_value == expected_value
+
+    # Default case
+    return False
+
+
 def match_feature_flag_properties(
-    flag, distinct_id, properties, cohort_properties=None
+    flag,
+    distinct_id,
+    properties,
+    cohort_properties=None,
+    flags_by_key=None,
+    evaluation_cache=None,
 ) -> FlagValue:
     flag_conditions = (flag.get("filters") or {}).get("groups") or []
     is_inconclusive = False
@@ -67,19 +220,18 @@ 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
             if is_condition_match(
-                flag, distinct_id, condition, properties, cohort_properties
+                flag,
+                distinct_id,
+                condition,
+                properties,
+                cohort_properties,
+                flags_by_key,
+                evaluation_cache,
             ):
                 variant_override = condition.get("variant")
                 if variant_override and variant_override in valid_variant_keys:
@@ -101,22 +253,36 @@ def match_feature_flag_properties(
 
 
 def is_condition_match(
-    feature_flag, distinct_id, condition, properties, cohort_properties
+    feature_flag,
+    distinct_id,
+    condition,
+    properties,
+    cohort_properties,
+    flags_by_key=None,
+    evaluation_cache=None,
 ) -> bool:
     rollout_percentage = condition.get("rollout_percentage")
     if len(condition.get("properties") or []) > 0:
         for prop in condition.get("properties"):
             property_type = prop.get("type")
             if property_type == "cohort":
-                matches = match_cohort(prop, properties, cohort_properties)
+                matches = match_cohort(
+                    prop,
+                    properties,
+                    cohort_properties,
+                    flags_by_key,
+                    evaluation_cache,
+                    distinct_id,
+                )
             elif property_type == "flag":
-                log.warning(
-                    "Flag dependency filters are not supported in local evaluation. "
-                    "Skipping condition for flag '%s' with dependency on flag '%s'",
-                    feature_flag.get("key", "unknown"),
-                    prop.get("key", "unknown"),
+                matches = evaluate_flag_dependency(
+                    prop,
+                    flags_by_key,
+                    evaluation_cache,
+                    distinct_id,
+                    properties,
+                    cohort_properties,
                 )
-                continue
             else:
                 matches = match_property(prop, properties)
             if not matches:
@@ -264,7 +430,14 @@ def match_property(property, property_values) -> bool:
     raise InconclusiveMatchError(f"Unknown operator {operator}")
 
 
-def match_cohort(property, property_values, cohort_properties) -> bool:
+def match_cohort(
+    property,
+    property_values,
+    cohort_properties,
+    flags_by_key=None,
+    evaluation_cache=None,
+    distinct_id=None,
+) -> bool:
     # Cohort properties are in the form of property groups like this:
     # {
     #     "cohort_id": {
@@ -281,10 +454,24 @@ def match_cohort(property, property_values, cohort_properties) -> bool:
         )
 
     property_group = cohort_properties[cohort_id]
-    return match_property_group(property_group, property_values, cohort_properties)
+    return match_property_group(
+        property_group,
+        property_values,
+        cohort_properties,
+        flags_by_key,
+        evaluation_cache,
+        distinct_id,
+    )
 
 
-def match_property_group(property_group, property_values, cohort_properties) -> bool:
+def match_property_group(
+    property_group,
+    property_values,
+    cohort_properties,
+    flags_by_key=None,
+    evaluation_cache=None,
+    distinct_id=None,
+) -> bool:
     if not property_group:
         return True
 
@@ -301,7 +488,14 @@ def match_property_group(property_group, property_values, cohort_properties) ->
         # a nested property group
         for prop in properties:
             try:
-                matches = match_property_group(prop, property_values, cohort_properties)
+                matches = match_property_group(
+                    prop,
+                    property_values,
+                    cohort_properties,
+                    flags_by_key,
+                    evaluation_cache,
+                    distinct_id,
+                )
                 if property_group_type == "AND":
                     if not matches:
                         return False
@@ -324,14 +518,23 @@ def match_property_group(property_group, property_values, cohort_properties) ->
         for prop in properties:
             try:
                 if prop.get("type") == "cohort":
-                    matches = match_cohort(prop, property_values, cohort_properties)
+                    matches = match_cohort(
+                        prop,
+                        property_values,
+                        cohort_properties,
+                        flags_by_key,
+                        evaluation_cache,
+                        distinct_id,
+                    )
                 elif prop.get("type") == "flag":
-                    log.warning(
-                        "Flag dependency filters are not supported in local evaluation. "
-                        "Skipping condition with dependency on flag '%s'",
-                        prop.get("key", "unknown"),
+                    matches = evaluate_flag_dependency(
+                        prop,
+                        flags_by_key,
+                        evaluation_cache,
+                        distinct_id,
+                        property_values,
+                        cohort_properties,
                     )
-                    continue
                 else:
                     matches = match_property(prop, property_values)
 

posthog/integrations/django.py

@@ -1,10 +1,18 @@
 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
+except ImportError:
+    # Fallback for older Django versions
+    import asyncio
+
+    iscoroutinefunction = asyncio.iscoroutinefunction
+
 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:
@@ -33,9 +41,24 @@ class PosthogContextMiddleware:
     frontend. See the documentation for `set_context_session` and `identify_context` for more details.
     """
 
+    # Django middleware capability flags
+    sync_capable = True
+    async_capable = True
+
     def __init__(self, get_response):
-        # type: (Callable[[HttpRequest], HttpResponse]) -> None
-        self.get_response = get_response
+        # type: (Union[Callable[[HttpRequest], HttpResponse], Callable[[HttpRequest], Awaitable[HttpResponse]]]) -> None
+        self._is_coroutine = iscoroutinefunction(get_response)
+        self._async_get_response = None  # type: Optional[Callable[[HttpRequest], Awaitable[HttpResponse]]]
+        self._sync_get_response = None  # type: Optional[Callable[[HttpRequest], HttpResponse]]
+
+        if self._is_coroutine:
+            self._async_get_response = cast(
+                "Callable[[HttpRequest], Awaitable[HttpResponse]]", get_response
+            )
+        else:
+            self._sync_get_response = cast(
+                "Callable[[HttpRequest], HttpResponse]", get_response
+            )
 
         from django.conf import settings
 
@@ -159,23 +182,39 @@ class PosthogContextMiddleware:
 
     def __call__(self, request):
         # type: (HttpRequest) -> HttpResponse
+        # Purely defensive around django's internal sync/async handling - this should be unreachable, but if it's reached, we may
+        # as well return something semi-meaningful
+        if self._is_coroutine:
+            raise RuntimeError(
+                "PosthogContextMiddleware received sync call but get_response is async"
+            )
+
         if self.request_filter and not self.request_filter(request):
-            return self.get_response(request)
+            assert self._sync_get_response is not None
+            return self._sync_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)
+            assert self._sync_get_response is not None
+            return self._sync_get_response(request)
 
-    def process_exception(self, request, exception):
+    async def __acall__(self, request):
+        # type: (HttpRequest) -> HttpResponse
         if self.request_filter and not self.request_filter(request):
-            return
+            if self._async_get_response is not None:
+                return await self._async_get_response(request)
+            else:
+                assert self._sync_get_response is not None
+                return self._sync_get_response(request)
 
-        if not self.capture_exceptions:
-            return
+        with contexts.new_context(self.capture_exceptions, client=self.client):
+            for k, v in self.extract_tags(request).items():
+                contexts.tag(k, v)
 
-        if self.client:
-            self.client.capture_exception(exception)
-        else:
-            capture_exception(exception)
+            if self._async_get_response is not None:
+                return await self._async_get_response(request)
+            else:
+                assert self._sync_get_response is not None
+                return self._sync_get_response(request)

posthog/request.py

@@ -132,12 +132,16 @@ def flags(
 
 
 def remote_config(
-    personal_api_key: str, host: Optional[str] = None, key: str = "", timeout: int = 15
+    personal_api_key: str,
+    project_api_key: str,
+    host: Optional[str] = None,
+    key: str = "",
+    timeout: int = 15,
 ) -> Any:
     """Get remote config flag value from remote_config API endpoint"""
     return get(
         personal_api_key,
-        f"/api/projects/@current/feature_flags/{key}/remote_config/",
+        f"/api/projects/@current/feature_flags/{key}/remote_config?token={project_api_key}",
         host,
         timeout,
     )

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"]
@@ -401,6 +517,78 @@ def test_new_client_override_defaults(
     assert props["urgent"] is True  # from call
 
 
+def test_vertex_ai_parameters_passed_through(
+    mock_client, mock_google_genai_client, mock_gemini_response
+):
+    """Test that Vertex AI parameters are properly passed to genai.Client"""
+    mock_google_genai_client.models.generate_content.return_value = mock_gemini_response
+
+    # Mock credentials object
+    mock_credentials = MagicMock()
+    mock_debug_config = MagicMock()
+    mock_http_options = MagicMock()
+
+    # Create client with Vertex AI parameters
+    Client(
+        vertexai=True,
+        credentials=mock_credentials,
+        project="test-project",
+        location="us-central1",
+        debug_config=mock_debug_config,
+        http_options=mock_http_options,
+        posthog_client=mock_client,
+    )
+
+    # Verify genai.Client was called with correct parameters
+    google_genai.Client.assert_called_once_with(
+        vertexai=True,
+        credentials=mock_credentials,
+        project="test-project",
+        location="us-central1",
+        debug_config=mock_debug_config,
+        http_options=mock_http_options,
+    )
+
+
+def test_api_key_mode(mock_client, mock_google_genai_client):
+    """Test API key authentication mode"""
+
+    # Create client with just API key (traditional mode)
+    Client(
+        api_key="test-api-key",
+        posthog_client=mock_client,
+    )
+
+    # Verify genai.Client was called with only api_key
+    google_genai.Client.assert_called_once_with(api_key="test-api-key")
+
+
+def test_vertex_ai_mode_with_optional_api_key(
+    mock_client, mock_google_genai_client, mock_gemini_response
+):
+    """Test Vertex AI mode with optional API key"""
+    mock_google_genai_client.models.generate_content.return_value = mock_gemini_response
+
+    mock_credentials = MagicMock()
+
+    # Create client with Vertex AI + API key
+    Client(
+        vertexai=True,
+        api_key="test-api-key",
+        credentials=mock_credentials,
+        project="test-project",
+        posthog_client=mock_client,
+    )
+
+    # Verify genai.Client was called with both Vertex AI params and API key
+    google_genai.Client.assert_called_once_with(
+        vertexai=True,
+        api_key="test-api-key",
+        credentials=mock_credentials,
+        project="test-project",
+    )
+
+
 def test_tool_use_response(mock_client, mock_google_genai_client, mock_gemini_response):
     """Test that tools defined in config are captured in $ai_tools property"""
     mock_google_genai_client.models.generate_content.return_value = mock_gemini_response
@@ -428,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",
@@ -557,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/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/test_client.py

@@ -2,17 +2,18 @@ import time
 import unittest
 from datetime import datetime
 from uuid import uuid4
-from posthog.contexts import get_context_session_id, set_context_session, new_context
 
 import mock
 import six
 from parameterized import parameterized
 
 from posthog.client import Client
+from posthog.contexts import get_context_session_id, new_context, set_context_session
 from posthog.request import APIError
 from posthog.test.test_utils import FAKE_TEST_API_KEY
 from posthog.types import FeatureFlag, LegacyFlagMetadata
 from posthog.version import VERSION
+from posthog.contexts import tag
 
 
 class TestClient(unittest.TestCase):
@@ -1741,6 +1742,7 @@ class TestClient(unittest.TestCase):
             person_properties={"distinct_id": "some_id"},
             group_properties={},
             geoip_disable=True,
+            flag_keys_to_evaluate=["random_key"],
         )
         patch_flags.reset_mock()
         client.feature_enabled(
@@ -1755,6 +1757,7 @@ class TestClient(unittest.TestCase):
             person_properties={"distinct_id": "feature_enabled_distinct_id"},
             group_properties={},
             geoip_disable=True,
+            flag_keys_to_evaluate=["random_key"],
         )
         patch_flags.reset_mock()
         client.get_all_flags_and_payloads("all_flags_payloads_id")
@@ -1815,6 +1818,7 @@ class TestClient(unittest.TestCase):
                 "instance": {"$group_key": "app.posthog.com"},
             },
             geoip_disable=False,
+            flag_keys_to_evaluate=["random_key"],
         )
 
         patch_flags.reset_mock()
@@ -1841,6 +1845,7 @@ class TestClient(unittest.TestCase):
                 "instance": {"$group_key": "app.posthog.com"},
             },
             geoip_disable=False,
+            flag_keys_to_evaluate=["random_key"],
         )
 
         patch_flags.reset_mock()
@@ -2057,7 +2062,7 @@ class TestClient(unittest.TestCase):
 
     def test_set_context_session_override_in_capture(self):
         """Test that explicit session ID overrides context session ID in capture"""
-        from posthog.contexts import set_context_session, new_context
+        from posthog.contexts import new_context, set_context_session
 
         with mock.patch("posthog.client.batch_post") as mock_post:
             client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True)
@@ -2158,6 +2163,7 @@ class TestClient(unittest.TestCase):
         self.assertEqual(result, {"test": "payload"})
         patch_remote_config.assert_called_once_with(
             "test-personal-key",
+            FAKE_TEST_API_KEY,
             client.host,
             "test-flag",
             timeout=client.feature_flags_request_timeout_seconds,
@@ -2185,6 +2191,7 @@ class TestClient(unittest.TestCase):
             "only_evaluate_locally": None,
             "person_properties": None,
             "group_properties": None,
+            "flag_keys_filter": None,
         }
         self.assertEqual(result, expected)
 
@@ -2195,6 +2202,7 @@ class TestClient(unittest.TestCase):
             "only_evaluate_locally": None,
             "person_properties": None,
             "group_properties": None,
+            "flag_keys_filter": None,
         }
         self.assertEqual(result, expected)
 
@@ -2210,6 +2218,7 @@ class TestClient(unittest.TestCase):
             "only_evaluate_locally": True,
             "person_properties": {"plan": "premium"},
             "group_properties": {"company": {"type": "enterprise"}},
+            "flag_keys_filter": None,
         }
         self.assertEqual(result, expected)
 
@@ -2221,6 +2230,7 @@ class TestClient(unittest.TestCase):
             "only_evaluate_locally": None,
             "person_properties": {"user_id": "123"},
             "group_properties": None,
+            "flag_keys_filter": None,
         }
         self.assertEqual(result, expected)
 
@@ -2231,6 +2241,7 @@ class TestClient(unittest.TestCase):
             "only_evaluate_locally": None,
             "person_properties": None,
             "group_properties": None,
+            "flag_keys_filter": None,
         }
         self.assertEqual(result, expected)
 
@@ -2247,6 +2258,53 @@ class TestClient(unittest.TestCase):
             client._parse_send_feature_flags(None)
         self.assertIn("Invalid type for send_feature_flags", str(cm.exception))
 
+    @mock.patch("posthog.client.flags")
+    def test_capture_with_send_feature_flags_flag_keys_filter(self, patch_flags):
+        """Test that SendFeatureFlagsOptions with flag_keys_filter only evaluates specified flags"""
+        # When flag_keys_to_evaluate is provided, the API should only return the requested flags
+        patch_flags.return_value = {
+            "featureFlags": {
+                "flag1": "value1",
+                "flag3": "value3",
+            }
+        }
+
+        with mock.patch("posthog.client.batch_post") as mock_post:
+            client = Client(
+                FAKE_TEST_API_KEY,
+                on_error=self.set_fail,
+                personal_api_key=FAKE_TEST_API_KEY,
+                sync_mode=True,
+            )
+
+            send_options = {
+                "flag_keys_filter": ["flag1", "flag3"],
+                "person_properties": {"subscription": "pro"},
+            }
+
+            msg_uuid = client.capture(
+                "test event", distinct_id="distinct_id", send_feature_flags=send_options
+            )
+
+            self.assertIsNotNone(msg_uuid)
+            self.assertFalse(self.failed)
+
+            # Verify flags() was called with flag_keys_to_evaluate
+            patch_flags.assert_called_once()
+            call_args = patch_flags.call_args[1]
+            self.assertEqual(call_args["flag_keys_to_evaluate"], ["flag1", "flag3"])
+            self.assertEqual(call_args["person_properties"], {"subscription": "pro"})
+
+            # Check the message includes only the filtered flags
+            mock_post.assert_called_once()
+            batch_data = mock_post.call_args[1]["batch"]
+            msg = batch_data[0]
+
+            self.assertEqual(msg["properties"]["$feature/flag1"], "value1")
+            self.assertEqual(msg["properties"]["$feature/flag3"], "value3")
+            # flag2 should not be included since it wasn't requested
+            self.assertNotIn("$feature/flag2", msg["properties"])
+
     @mock.patch("posthog.client.batch_post")
     def test_get_feature_flag_result_with_empty_string_payload(self, patch_batch_post):
         """Test that get_feature_flag_result returns a FeatureFlagResult when payload is empty string"""
@@ -2282,9 +2340,7 @@ class TestClient(unittest.TestCase):
                             }
                         ]
                     },
-                    "payloads": {
-                        "empty-variant": ""  # Empty string payload
-                    },
+                    "payloads": {"empty-variant": ""},  # Empty string payload
                 },
             }
         ]
@@ -2355,3 +2411,58 @@ class TestClient(unittest.TestCase):
         self.assertEqual(
             result["featureFlagPayloads"]["normal-payload-flag"], "normal payload"
         )
+
+    def test_context_tags_added(self):
+        with mock.patch("posthog.client.batch_post") as mock_post:
+            client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True)
+
+            with new_context():
+                tag("random_tag", 12345)
+                client.capture("python test event", distinct_id="distinct_id")
+
+            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_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/types.py

@@ -9,6 +9,7 @@ FlagValue = Union[bool, str]
 BeforeSendCallback = Callable[[dict[str, Any]], Optional[dict[str, Any]]]
 
 
+# Type alias for the send_feature_flags parameter
 class SendFeatureFlagsOptions(TypedDict, total=False):
     """Options for sending feature flags with capture events.
 
@@ -22,9 +23,11 @@ class SendFeatureFlagsOptions(TypedDict, total=False):
             Format: { group_type_name: { group_properties } }
     """
 
+    should_send: bool
     only_evaluate_locally: Optional[bool]
     person_properties: Optional[dict[str, Any]]
     group_properties: Optional[dict[str, dict[str, Any]]]
+    flag_keys_filter: Optional[list[str]]
 
 
 @dataclass(frozen=True)

posthog/version.py

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

remote_config_example.py

@@ -0,0 +1,32 @@
+#!/usr/bin/env python3
+"""
+Simple test script for PostHog remote config endpoint.
+"""
+
+import posthog
+
+# Initialize PostHog client
+posthog.api_key = "phc_..."
+posthog.personal_api_key = "phs_..."  # or "phx_..."
+posthog.host = "http://localhost:8000"  # or "https://us.posthog.com"
+posthog.debug = True
+
+
+def test_remote_config():
+    """Test remote config payload retrieval."""
+    print("Testing remote config endpoint...")
+
+    # Test feature flag key - replace with an actual flag key from your project
+    flag_key = "unencrypted-remote-config-setting"
+
+    try:
+        # Get remote config payload
+        payload = posthog.get_remote_config_payload(flag_key)
+        print(f"✅ Success! Remote config payload for '{flag_key}': {payload}")
+
+    except Exception as e:
+        print(f"❌ Error getting remote config: {e}")
+
+
+if __name__ == "__main__":
+    test_remote_config()