fix(llma): gemini missing cached and reasoning tokens (#323)

* fix(llma): Gemini missing cached and reasoning tokens

* chore(llma): bump version

* chore(llma): run ruff

.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

CHANGELOG.md

@@ -1,4 +1,57 @@
-# 6.2.1 - 2025-06-21
+# 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
+
+# 6.3.3 - 2025-08-01
+
+- fix: `get_feature_flag_result` now correctly returns FeatureFlagResult when payload is empty string instead of None
+
+# 6.3.2 - 2025-07-31
+
+- fix: Anthropic's tool calls are now handled properly
+
+# 6.3.0 - 2025-07-22
+
+- feat: Enhanced `send_feature_flags` parameter to accept `SendFeatureFlagsOptions` object for declarative control over local/remote evaluation and custom properties
+
+# 6.2.1 - 2025-07-21
 
 - feat: make `posthog_client` an optional argument in PostHog AI providers wrappers (`posthog.ai.*`), intuitively using the default client as the default
 

bin/docs_scripts/doc_constant.py

@@ -4,49 +4,9 @@ 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",
-]
-
 # Documentation generation metadata
 DOCUMENTATION_METADATA = {
-    "hogRef": "0.1",
+    "hogRef": "0.3",
     "slugPrefix": "posthog-python",
     "specUrl": "https://github.com/PostHog/posthog-python",
 }

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
@@ -459,12 +480,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,175 +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"))
-exit()
-# # 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/__init__.py

@@ -11,7 +11,7 @@ from posthog.contexts import (
     set_context_session as inner_set_context_session,
     identify_context as inner_identify_context,
 )
-from posthog.types import FeatureFlag, FlagsAndPayloads
+from posthog.types import FeatureFlag, FlagsAndPayloads, FeatureFlagResult
 from posthog.version import VERSION
 
 __version__ = VERSION
@@ -388,9 +388,9 @@ def capture_exception(
 def feature_enabled(
     key,  # type: str
     distinct_id,  # type: str
-    groups={},  # type: dict
-    person_properties={},  # type: dict
-    group_properties={},  # type: dict
+    groups=None,  # type: Optional[dict]
+    person_properties=None,  # type: Optional[dict]
+    group_properties=None,  # type: Optional[dict]
     only_evaluate_locally=False,  # type: bool
     send_feature_flag_events=True,  # type: bool
     disable_geoip=None,  # type: Optional[bool]
@@ -427,9 +427,9 @@ def feature_enabled(
         "feature_enabled",
         key=key,
         distinct_id=distinct_id,
-        groups=groups,
-        person_properties=person_properties,
-        group_properties=group_properties,
+        groups=groups or {},
+        person_properties=person_properties or {},
+        group_properties=group_properties or {},
         only_evaluate_locally=only_evaluate_locally,
         send_feature_flag_events=send_feature_flag_events,
         disable_geoip=disable_geoip,
@@ -439,9 +439,9 @@ def feature_enabled(
 def get_feature_flag(
     key,  # type: str
     distinct_id,  # type: str
-    groups={},  # type: dict
-    person_properties={},  # type: dict
-    group_properties={},  # type: dict
+    groups=None,  # type: Optional[dict]
+    person_properties=None,  # type: Optional[dict]
+    group_properties=None,  # type: Optional[dict]
     only_evaluate_locally=False,  # type: bool
     send_feature_flag_events=True,  # type: bool
     disable_geoip=None,  # type: Optional[bool]
@@ -477,9 +477,9 @@ def get_feature_flag(
         "get_feature_flag",
         key=key,
         distinct_id=distinct_id,
-        groups=groups,
-        person_properties=person_properties,
-        group_properties=group_properties,
+        groups=groups or {},
+        person_properties=person_properties or {},
+        group_properties=group_properties or {},
         only_evaluate_locally=only_evaluate_locally,
         send_feature_flag_events=send_feature_flag_events,
         disable_geoip=disable_geoip,
@@ -488,9 +488,9 @@ def get_feature_flag(
 
 def get_all_flags(
     distinct_id,  # type: str
-    groups={},  # type: dict
-    person_properties={},  # type: dict
-    group_properties={},  # type: dict
+    groups=None,  # type: Optional[dict]
+    person_properties=None,  # type: Optional[dict]
+    group_properties=None,  # type: Optional[dict]
     only_evaluate_locally=False,  # type: bool
     disable_geoip=None,  # type: Optional[bool]
 ) -> Optional[dict[str, FeatureFlag]]:
@@ -520,21 +520,64 @@ def get_all_flags(
     return _proxy(
         "get_all_flags",
         distinct_id=distinct_id,
-        groups=groups,
-        person_properties=person_properties,
-        group_properties=group_properties,
+        groups=groups or {},
+        person_properties=person_properties or {},
+        group_properties=group_properties or {},
         only_evaluate_locally=only_evaluate_locally,
         disable_geoip=disable_geoip,
     )
 
 
+def get_feature_flag_result(
+    key,
+    distinct_id,
+    groups=None,  # type: Optional[dict]
+    person_properties=None,  # type: Optional[dict]
+    group_properties=None,  # type: Optional[dict]
+    only_evaluate_locally=False,
+    send_feature_flag_events=True,
+    disable_geoip=None,  # type: Optional[bool]
+):
+    # type: (...) -> Optional[FeatureFlagResult]
+    """
+    Get a FeatureFlagResult object which contains the flag result and payload.
+
+    This method evaluates a feature flag and returns a FeatureFlagResult object containing:
+    - enabled: Whether the flag is enabled
+    - variant: The variant value if the flag has variants
+    - payload: The payload associated with the flag (automatically deserialized from JSON)
+    - key: The flag key
+    - reason: Why the flag was enabled/disabled
+
+    Example:
+    ```python
+    result = posthog.get_feature_flag_result('beta-feature', 'distinct_id')
+    if result and result.enabled:
+        # Use the variant and payload
+        print(f"Variant: {result.variant}")
+        print(f"Payload: {result.payload}")
+    ```
+    """
+    return _proxy(
+        "get_feature_flag_result",
+        key=key,
+        distinct_id=distinct_id,
+        groups=groups or {},
+        person_properties=person_properties or {},
+        group_properties=group_properties or {},
+        only_evaluate_locally=only_evaluate_locally,
+        send_feature_flag_events=send_feature_flag_events,
+        disable_geoip=disable_geoip,
+    )
+
+
 def get_feature_flag_payload(
     key,
     distinct_id,
     match_value=None,
-    groups={},
-    person_properties={},
-    group_properties={},
+    groups=None,  # type: Optional[dict]
+    person_properties=None,  # type: Optional[dict]
+    group_properties=None,  # type: Optional[dict]
     only_evaluate_locally=False,
     send_feature_flag_events=True,
     disable_geoip=None,  # type: Optional[bool]
@@ -544,9 +587,9 @@ def get_feature_flag_payload(
         key=key,
         distinct_id=distinct_id,
         match_value=match_value,
-        groups=groups,
-        person_properties=person_properties,
-        group_properties=group_properties,
+        groups=groups or {},
+        person_properties=person_properties or {},
+        group_properties=group_properties or {},
         only_evaluate_locally=only_evaluate_locally,
         send_feature_flag_events=send_feature_flag_events,
         disable_geoip=disable_geoip,
@@ -575,18 +618,18 @@ def get_remote_config_payload(
 
 def get_all_flags_and_payloads(
     distinct_id,
-    groups={},
-    person_properties={},
-    group_properties={},
+    groups=None,  # type: Optional[dict]
+    person_properties=None,  # type: Optional[dict]
+    group_properties=None,  # type: Optional[dict]
     only_evaluate_locally=False,
     disable_geoip=None,  # type: Optional[bool]
 ) -> FlagsAndPayloads:
     return _proxy(
         "get_all_flags_and_payloads",
         distinct_id=distinct_id,
-        groups=groups,
-        person_properties=person_properties,
-        group_properties=group_properties,
+        groups=groups or {},
+        person_properties=person_properties or {},
+        group_properties=group_properties or {},
         only_evaluate_locally=only_evaluate_locally,
         disable_geoip=disable_geoip,
     )

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,8 @@ import time
 import uuid
 from typing import Any, Dict, Optional
 
+from posthog.ai.types import TokenUsage
+
 try:
     from google import genai
 except ImportError:
@@ -13,9 +15,16 @@ 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 (
+    format_gemini_input,
+    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())
+        from posthog.ai.types import StreamingEventData
 
-        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 {}),
-        }
+        # Prepare standardized event data
+        formatted_input = self._format_input(contents)
+        sanitized_input = sanitize_gemini(formatted_input)
 
-        if distinct_id is None:
-            event_properties["$process_person_profile"] = False
+        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 hasattr(self._ph_client, "capture"):
-            self._ph_client.capture(
-                distinct_id=distinct_id,
-                event="$ai_generation",
-                properties=event_properties,
-                groups=groups,
-            )
+        # Use the common capture function
+        capture_streaming_event(self._ph_client, event_data)
 
     def _format_input(self, contents):
         """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)}]
+
+        return format_gemini_input(contents)
 
     def generate_content_stream(
         self,

posthog/ai/gemini/gemini_converter.py

@@ -0,0 +1,460 @@
+"""
+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_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(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

@@ -5,6 +5,7 @@ except ImportError:
         "Please install LangChain to use this feature: 'pip install langchain'"
     )
 
+import json
 import logging
 import time
 from dataclasses import dataclass
@@ -29,12 +30,14 @@ from langchain_core.messages import (
     HumanMessage,
     SystemMessage,
     ToolMessage,
+    ToolCall,
 )
 from langchain_core.outputs import ChatGeneration, LLMResult
 from pydantic import BaseModel
 
-from posthog import default_client
+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")
@@ -81,7 +84,7 @@ class CallbackHandler(BaseCallbackHandler):
     The PostHog LLM observability callback handler for LangChain.
     """
 
-    _client: Client
+    _ph_client: Client
     """PostHog client instance."""
 
     _distinct_id: Optional[Union[str, int, UUID]]
@@ -127,10 +130,7 @@ class CallbackHandler(BaseCallbackHandler):
             privacy_mode: Whether to redact the input and output of the trace.
             groups: Optional additional PostHog groups to use for the trace.
         """
-        posthog_client = client or default_client
-        if posthog_client is None:
-            raise ValueError("PostHog client is required")
-        self._client = posthog_client
+        self._ph_client = client or setup()
         self._distinct_id = distinct_id
         self._trace_id = trace_id
         self._properties = properties or {}
@@ -481,7 +481,7 @@ class CallbackHandler(BaseCallbackHandler):
         event_properties = {
             "$ai_trace_id": trace_id,
             "$ai_input_state": with_privacy_mode(
-                self._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,
@@ -497,13 +497,13 @@ class CallbackHandler(BaseCallbackHandler):
             event_properties["$ai_is_error"] = True
         elif outputs is not None:
             event_properties["$ai_output_state"] = with_privacy_mode(
-                self._client, self._privacy_mode, outputs
+                self._ph_client, self._privacy_mode, outputs
             )
 
         if self._distinct_id is None:
             event_properties["$process_person_profile"] = False
 
-        self._client.capture(
+        self._ph_client.capture(
             distinct_id=self._distinct_id or run_id,
             event=event_name,
             properties=event_properties,
@@ -550,17 +550,16 @@ class CallbackHandler(BaseCallbackHandler):
             "$ai_provider": run.provider,
             "$ai_model": run.model,
             "$ai_model_parameters": run.model_params,
-            "$ai_input": with_privacy_mode(self._client, self._privacy_mode, run.input),
+            "$ai_input": with_privacy_mode(
+                self._ph_client, self._privacy_mode, sanitize_langchain(run.input)
+            ),
             "$ai_http_status": 200,
             "$ai_latency": run.latency,
             "$ai_base_url": run.base_url,
         }
+
         if run.tools:
-            event_properties["$ai_tools"] = with_privacy_mode(
-                self._client,
-                self._privacy_mode,
-                run.tools,
-            )
+            event_properties["$ai_tools"] = run.tools
 
         if isinstance(output, BaseException):
             event_properties["$ai_http_status"] = _get_http_status(output)
@@ -586,10 +585,11 @@ class CallbackHandler(BaseCallbackHandler):
                 ]
             else:
                 completions = [
-                    _extract_raw_esponse(generation) for generation in generation_result
+                    _extract_raw_response(generation)
+                    for generation in generation_result
                 ]
             event_properties["$ai_output_choices"] = with_privacy_mode(
-                self._client, self._privacy_mode, completions
+                self._ph_client, self._privacy_mode, completions
             )
 
         if self._properties:
@@ -598,7 +598,7 @@ class CallbackHandler(BaseCallbackHandler):
         if self._distinct_id is None:
             event_properties["$process_person_profile"] = False
 
-        self._client.capture(
+        self._ph_client.capture(
             distinct_id=self._distinct_id or trace_id,
             event="$ai_generation",
             properties=event_properties,
@@ -617,7 +617,7 @@ class CallbackHandler(BaseCallbackHandler):
         )
 
 
-def _extract_raw_esponse(last_response):
+def _extract_raw_response(last_response):
     """Extract the response from the last response of the LLM call."""
     # We return the text of the response if not empty
     if last_response.text is not None and last_response.text.strip() != "":
@@ -630,12 +630,35 @@ def _extract_raw_esponse(last_response):
         return ""
 
 
-def _convert_message_to_dict(message: BaseMessage) -> Dict[str, Any]:
+def _convert_lc_tool_calls_to_oai(
+    tool_calls: list[ToolCall],
+) -> list[dict[str, Any]]:
+    try:
+        return [
+            {
+                "type": "function",
+                "id": tool_call["id"],
+                "function": {
+                    "name": tool_call["name"],
+                    "arguments": json.dumps(tool_call["args"]),
+                },
+            }
+            for tool_call in tool_calls
+        ]
+    except KeyError:
+        return tool_calls
+
+
+def _convert_message_to_dict(message: BaseMessage) -> dict[str, Any]:
     # assistant message
     if isinstance(message, HumanMessage):
         message_dict = {"role": "user", "content": message.content}
     elif isinstance(message, AIMessage):
         message_dict = {"role": "assistant", "content": message.content}
+        if message.tool_calls:
+            message_dict["tool_calls"] = _convert_lc_tool_calls_to_oai(
+                message.tool_calls
+            )
     elif isinstance(message, SystemMessage):
         message_dict = {"role": "system", "content": message.content}
     elif isinstance(message, ToolMessage):
@@ -648,6 +671,9 @@ def _convert_message_to_dict(message: BaseMessage) -> Dict[str, Any]:
     if message.additional_kwargs:
         message_dict.update(message.additional_kwargs)
 
+    if "content" in message_dict and not message_dict["content"]:
+        message_dict["content"] = ""
+
     return message_dict
 
 

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:
@@ -11,9 +13,17 @@ except ImportError:
 
 from posthog.ai.utils import (
     call_llm_and_track_usage,
-    get_model_params,
+    extract_available_tool_calls,
+    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
 
@@ -32,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()
 
@@ -111,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)
 
@@ -121,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
 
@@ -167,6 +160,7 @@ class WrappedResponses:
                     usage_stats,
                     latency,
                     output,
+                    None,  # Responses API doesn't have tools
                 )
 
         return generator()
@@ -179,56 +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,
-        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("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 tool_calls:
-            event_properties["$ai_tools"] = with_privacy_mode(
-                self._client._ph_client,
-                posthog_privacy_mode,
-                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,
@@ -339,9 +317,9 @@ class WrappedCompletions:
         **kwargs: Any,
     ):
         start_time = time.time()
-        usage_stats: Dict[str, int] = {}
+        usage_stats: TokenUsage = TokenUsage()
         accumulated_content = []
-        accumulated_tools = {}
+        accumulated_tool_calls: Dict[int, Dict[str, Any]] = {}
         if "stream_options" not in kwargs:
             kwargs["stream_options"] = {}
         kwargs["stream_options"]["include_usage"] = True
@@ -350,70 +328,42 @@ class WrappedCompletions:
         def generator():
             nonlocal usage_stats
             nonlocal accumulated_content  # noqa: F824
-            nonlocal accumulated_tools  # 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)
 
-                        # Process tool calls
-                        tool_calls = getattr(chunk.choices[0].delta, "tool_calls", None)
-                        if tool_calls:
-                            for tool_call in tool_calls:
-                                index = tool_call.index
-                                if index not in accumulated_tools:
-                                    accumulated_tools[index] = tool_call
-                                else:
-                                    # Append arguments for existing tool calls
-                                    if hasattr(tool_call, "function") and hasattr(
-                                        tool_call.function, "arguments"
-                                    ):
-                                        accumulated_tools[
-                                            index
-                                        ].function.arguments += (
-                                            tool_call.function.arguments
-                                        )
+                    # 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)
-                tools = list(accumulated_tools.values()) if accumulated_tools else None
+
+                # 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,
@@ -423,8 +373,9 @@ class WrappedCompletions:
                     kwargs,
                     usage_stats,
                     latency,
-                    output,
-                    tools,
+                    accumulated_content,
+                    tool_calls_list,
+                    extract_available_tool_calls("openai", kwargs),
                 )
 
         return generator()
@@ -437,56 +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 tool_calls:
-            event_properties["$ai_tools"] = with_privacy_mode(
-                self._client._ph_client,
-                posthog_privacy_mode,
-                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:
@@ -523,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())
 
@@ -545,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
@@ -12,9 +14,19 @@ except ImportError:
 from posthog import setup
 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
 
 
@@ -33,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()
 
@@ -65,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(
@@ -112,9 +126,9 @@ 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)
+        response = self._original.create(**kwargs)
 
         async def async_generator():
             nonlocal usage_stats
@@ -122,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
 
@@ -158,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,
@@ -168,6 +165,7 @@ class WrappedResponses:
                     usage_stats,
                     latency,
                     output,
+                    extract_available_tool_calls("openai", kwargs),
                 )
 
         return async_generator()
@@ -180,10 +178,10 @@ 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,
-        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())
@@ -193,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),
@@ -213,12 +213,8 @@ class WrappedResponses:
             **(posthog_properties or {}),
         }
 
-        if tool_calls:
-            event_properties["$ai_tools"] = with_privacy_mode(
-                self._client._ph_client,
-                posthog_privacy_mode,
-                tool_calls,
-            )
+        if available_tool_calls:
+            event_properties["$ai_tools"] = available_tool_calls
 
         if posthog_distinct_id is None:
             event_properties["$process_person_profile"] = False
@@ -342,82 +338,52 @@ class WrappedCompletions:
         **kwargs: Any,
     ):
         start_time = time.time()
-        usage_stats: Dict[str, int] = {}
+        usage_stats: TokenUsage = TokenUsage()
         accumulated_content = []
-        accumulated_tools = {}
+        accumulated_tool_calls: Dict[int, Dict[str, Any]] = {}
 
         if "stream_options" not in kwargs:
             kwargs["stream_options"] = {}
         kwargs["stream_options"]["include_usage"] = True
-        response = await self._original.create(**kwargs)
+        response = self._original.create(**kwargs)
 
         async def async_generator():
             nonlocal usage_stats
             nonlocal accumulated_content  # noqa: F824
-            nonlocal accumulated_tools  # 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)
-
-                        # Process tool calls
-                        tool_calls = getattr(chunk.choices[0].delta, "tool_calls", None)
-                        if tool_calls:
-                            for tool_call in tool_calls:
-                                index = tool_call.index
-                                if index not in accumulated_tools:
-                                    accumulated_tools[index] = tool_call
-                                else:
-                                    # Append arguments for existing tool calls
-                                    if hasattr(tool_call, "function") and hasattr(
-                                        tool_call.function, "arguments"
-                                    ):
-                                        accumulated_tools[
-                                            index
-                                        ].function.arguments += (
-                                            tool_call.function.arguments
-                                        )
+                    # 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)
-                tools = list(accumulated_tools.values()) if accumulated_tools else None
+
+                # 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,
@@ -427,8 +393,9 @@ class WrappedCompletions:
                     kwargs,
                     usage_stats,
                     latency,
-                    output,
-                    tools,
+                    accumulated_content,
+                    tool_calls_list,
+                    extract_available_tool_calls("openai", kwargs),
                 )
 
         return async_generator()
@@ -441,10 +408,11 @@ 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())
@@ -454,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
             ),
@@ -474,12 +444,8 @@ class WrappedCompletions:
             **(posthog_properties or {}),
         }
 
-        if tool_calls:
-            event_properties["$ai_tools"] = with_privacy_mode(
-                self._client._ph_client,
-                posthog_privacy_mode,
-                tool_calls,
-            )
+        if available_tool_calls:
+            event_properties["$ai_tools"] = available_tool_calls
 
         if posthog_distinct_id is None:
             event_properties["$process_person_profile"] = False
@@ -502,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(
@@ -527,20 +494,22 @@ class WrappedEmbeddings:
         Returns:
             The response from OpenAI's embeddings.create call.
         """
+
         if posthog_trace_id is None:
             posthog_trace_id = str(uuid.uuid4())
 
         start_time = time.time()
-        response = await self._original.create(**kwargs)
+        response = self._original.create(**kwargs)
         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
 
@@ -549,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),
@@ -583,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
@@ -599,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
@@ -615,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,612 @@
+"""
+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
+    """
+    if api_type == "chat":
+        return kwargs.get("messages")
+    else:  # responses API
+        return kwargs.get("input")

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, Optional
 
 from posthog.client import Client as PostHogClient
+from posthog.ai.types import 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,285 +93,128 @@ 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
+
+        return format_gemini_response(response)
+    return []
 
 
-def format_response_anthropic(response):
-    output = []
-    for choice in response.content:
-        if choice.text:
-            output.append(
-                {
-                    "role": "assistant",
-                    "content": choice.text,
-                }
-            )
-    return output
-
-
-def format_response_openai(response):
-    output = []
-    if hasattr(response, "choices"):
-        for choice in response.choices:
-            # Handle Chat Completions response format
-            if hasattr(choice, "message") and choice.message and choice.message.content:
-                output.append(
-                    {
-                        "content": choice.message.content,
-                        "role": choice.message.role,
-                    }
-                )
-    # Handle Responses API format
-    if hasattr(response, "output"):
-        for item in response.output:
-            if item.type == "message":
-                # Extract text content from the content list
-                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")
-                        ):
-                            output.append(
-                                {
-                                    "content": content_item.text,
-                                    "role": item.role,
-                                }
-                            )
-                        elif hasattr(content_item, "text"):
-                            output.append(
-                                {
-                                    "content": content_item.text,
-                                    "role": item.role,
-                                }
-                            )
-                        elif (
-                            hasattr(content_item, "type")
-                            and content_item.type == "input_image"
-                            and hasattr(content_item, "image_url")
-                        ):
-                            output.append(
-                                {
-                                    "content": {
-                                        "type": "image",
-                                        "image": content_item.image_url,
-                                    },
-                                    "role": item.role,
-                                }
-                            )
-                else:
-                    output.append(
-                        {
-                            "content": item.content,
-                            "role": item.role,
-                        }
-                    )
-    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_text = ""
-                if hasattr(candidate.content, "parts") and candidate.content.parts:
-                    for part in candidate.content.parts:
-                        if hasattr(part, "text") and part.text:
-                            content_text += part.text
-                if content_text:
-                    output.append(
-                        {
-                            "role": "assistant",
-                            "content": content_text,
-                        }
-                    )
-            elif hasattr(candidate, "text") and candidate.text:
-                output.append(
-                    {
-                        "role": "assistant",
-                        "content": candidate.text,
-                    }
-                )
-    elif hasattr(response, "text") and response.text:
-        output.append(
-            {
-                "role": "assistant",
-                "content": response.text,
-            }
-        )
-    return output
-
-
-def format_tool_calls(response, provider: str):
+def extract_available_tool_calls(provider: str, kwargs: Dict[str, Any]):
+    """
+    Extract available tool calls for the given provider.
+    """
     if provider == "anthropic":
-        if hasattr(response, "tools") and response.tools and len(response.tools) > 0:
-            return response.tools
-    elif provider == "openai":
-        # Handle both Chat Completions and Responses API
-        if hasattr(response, "choices") and response.choices:
-            # Check for tool_calls in message (Chat Completions format)
-            if (
-                hasattr(response.choices[0], "message")
-                and hasattr(response.choices[0].message, "tool_calls")
-                and response.choices[0].message.tool_calls
-            ):
-                return response.choices[0].message.tool_calls
+        from posthog.ai.anthropic.anthropic_converter import extract_anthropic_tools
 
-            # Check for tool_calls directly in response (Responses API format)
-            if (
-                hasattr(response.choices[0], "tool_calls")
-                and response.choices[0].tool_calls
-            ):
-                return response.choices[0].tool_calls
+        return extract_anthropic_tools(kwargs)
+    elif provider == "gemini":
+        from posthog.ai.gemini.gemini_converter import extract_gemini_tools
+
+        return extract_gemini_tools(kwargs)
+    elif provider == "openai":
+        from posthog.ai.openai.openai_converter import extract_openai_tools
+
+        return extract_openai_tools(kwargs)
     return None
 
 
 def merge_system_prompt(kwargs: Dict[str, Any], provider: str):
-    messages: List[Dict[str, Any]] = []
+    """
+    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
+
         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)}]
+        return format_gemini_input(contents)
+    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:
+                messages = [
+                    {"role": "system", "content": kwargs.get("system")}
+                ] + 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
+                messages = [
+                    {"role": "system", "content": kwargs.get("instructions")}
+                ] + messages
+
+        return messages
+
+    # Default case - return empty list
+    return []
 
 
 def call_llm_and_track_usage(
@@ -318,7 +225,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:
@@ -330,8 +237,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)
@@ -358,12 +265,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)
             ),
@@ -377,33 +287,22 @@ def call_llm_and_track_usage(
             **(error_params or {}),
         }
 
-        tool_calls = format_tool_calls(response, provider)
-        if tool_calls:
-            event_properties["$ai_tools"] = with_privacy_mode(
-                ph_client, posthog_privacy_mode, tool_calls
-            )
+        available_tool_calls = extract_available_tool_calls(provider, kwargs)
 
-        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
-            )
+        if available_tool_calls:
+            event_properties["$ai_tools"] = available_tool_calls
 
-        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_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("reasoning_tokens") is not None
-            and usage.get("reasoning_tokens", 0) > 0
-        ):
-            event_properties["$ai_reasoning_tokens"] = usage.get("reasoning_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
+
+        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
@@ -437,7 +336,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:
@@ -445,8 +344,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)
@@ -473,12 +372,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)
             ),
@@ -492,27 +394,18 @@ async def call_llm_and_track_usage_async(
             **(error_params or {}),
         }
 
-        tool_calls = format_tool_calls(response, provider)
-        if tool_calls:
-            event_properties["$ai_tools"] = with_privacy_mode(
-                ph_client, posthog_privacy_mode, tool_calls
-            )
+        available_tool_calls = extract_available_tool_calls(provider, kwargs)
 
-        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
-            )
+        if available_tool_calls:
+            event_properties["$ai_tools"] = available_tool_calls
 
-        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_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
+
+        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
@@ -538,7 +431,122 @@ async def call_llm_and_track_usage_async(
     return response
 
 
+def sanitize_messages(data: Any, provider: str) -> Any:
+    """Sanitize messages using provider-specific sanitization functions."""
+    if provider == "anthropic":
+        return sanitize_anthropic(data)
+    elif provider == "openai":
+        return sanitize_openai(data)
+    elif provider == "gemini":
+        return sanitize_gemini(data)
+    elif provider == "langchain":
+        return sanitize_langchain(data)
+    return data
+
+
 def with_privacy_mode(ph_client: PostHogClient, privacy_mode: bool, value: Any):
     if ph_client.privacy_mode or privacy_mode:
         return None
     return value
+
+
+def capture_streaming_event(
+    ph_client: PostHogClient,
+    event_data: StreamingEventData,
+):
+    """
+    Unified streaming event capture for all LLM providers.
+
+    This function handles the common logic for capturing streaming events across all providers.
+    All provider-specific formatting should be done BEFORE calling this function.
+
+    The function handles:
+    - Building PostHog event properties
+    - Extracting and adding tools based on provider
+    - Applying privacy mode
+    - Adding special token fields (cache, reasoning)
+    - Provider-specific fields (e.g., OpenAI instructions)
+    - Sending the event to PostHog
+
+    Args:
+        ph_client: PostHog client instance
+        event_data: Standardized streaming event data containing all necessary information
+    """
+    trace_id = event_data.get("trace_id") or str(uuid.uuid4())
+
+    # Build base event properties
+    event_properties = {
+        "$ai_provider": event_data["provider"],
+        "$ai_model": event_data["model"],
+        "$ai_model_parameters": get_model_params(event_data["kwargs"]),
+        "$ai_input": with_privacy_mode(
+            ph_client,
+            event_data["privacy_mode"],
+            event_data["formatted_input"],
+        ),
+        "$ai_output_choices": with_privacy_mode(
+            ph_client,
+            event_data["privacy_mode"],
+            event_data["formatted_output"],
+        ),
+        "$ai_http_status": 200,
+        "$ai_input_tokens": event_data["usage_stats"].get("input_tokens", 0),
+        "$ai_output_tokens": event_data["usage_stats"].get("output_tokens", 0),
+        "$ai_latency": event_data["latency"],
+        "$ai_trace_id": trace_id,
+        "$ai_base_url": str(event_data["base_url"]),
+        **(event_data.get("properties") or {}),
+    }
+
+    # Extract and add tools based on provider
+    available_tools = extract_available_tool_calls(
+        event_data["provider"],
+        event_data["kwargs"],
+    )
+    if available_tools:
+        event_properties["$ai_tools"] = available_tools
+
+    # Add optional token fields
+    # For Anthropic, always include cache fields even if 0 (backward compatibility)
+    # For others, only include if present and non-zero
+    if event_data["provider"] == "anthropic":
+        # Anthropic always includes cache fields
+        cache_read = event_data["usage_stats"].get("cache_read_input_tokens", 0)
+        cache_creation = event_data["usage_stats"].get("cache_creation_input_tokens", 0)
+        event_properties["$ai_cache_read_input_tokens"] = cache_read
+        event_properties["$ai_cache_creation_input_tokens"] = cache_creation
+    else:
+        # Other providers only include if non-zero
+        optional_token_fields = [
+            "cache_read_input_tokens",
+            "cache_creation_input_tokens",
+            "reasoning_tokens",
+        ]
+
+        for field in optional_token_fields:
+            value = event_data["usage_stats"].get(field)
+            if value is not None and isinstance(value, int) and value > 0:
+                event_properties[f"$ai_{field}"] = value
+
+    # Handle provider-specific fields
+    if (
+        event_data["provider"] == "openai"
+        and event_data["kwargs"].get("instructions") is not None
+    ):
+        event_properties["$ai_instructions"] = with_privacy_mode(
+            ph_client,
+            event_data["privacy_mode"],
+            event_data["kwargs"]["instructions"],
+        )
+
+    if event_data.get("distinct_id") is None:
+        event_properties["$process_person_profile"] = False
+
+    # Send event to PostHog
+    if hasattr(ph_client, "capture"):
+        ph_client.capture(
+            distinct_id=event_data.get("distinct_id") or trace_id,
+            event="$ai_generation",
+            properties=event_properties,
+            groups=event_data.get("groups"),
+        )

posthog/args.py

@@ -5,6 +5,8 @@ from datetime import datetime
 import numbers
 from uuid import UUID
 
+from posthog.types import SendFeatureFlagsOptions
+
 ID_TYPES = Union[numbers.Number, str, UUID, int]
 
 
@@ -22,7 +24,8 @@ class OptionalCaptureArgs(TypedDict):
             error ID if you capture an exception).
         groups: Group identifiers to associate with this event (format: {group_type: group_key})
         send_feature_flags: Whether to include currently active feature flags in the event properties.
-            Defaults to False
+            Can be a boolean (True/False) or a SendFeatureFlagsOptions object for advanced configuration.
+            Defaults to False.
         disable_geoip: Whether to disable GeoIP lookup for this event. Defaults to False.
     """
 
@@ -32,8 +35,8 @@ class OptionalCaptureArgs(TypedDict):
     uuid: NotRequired[Optional[str]]
     groups: NotRequired[Optional[Dict[str, str]]]
     send_feature_flags: NotRequired[
-        Optional[bool]
-    ]  # Optional so we can tell if the user is intentionally overriding a client setting or not
+        Optional[Union[bool, SendFeatureFlagsOptions]]
+    ]  # Updated to support both boolean and options object
     disable_geoip: NotRequired[
         Optional[bool]
     ]  # As above, optional so we can tell if the user is intentionally overriding a client setting or not

posthog/client.py

@@ -44,6 +44,7 @@ from posthog.types import (
     FlagsAndPayloads,
     FlagsResponse,
     FlagValue,
+    SendFeatureFlagsOptions,
     normalize_flags_response,
     to_flags_and_payloads,
     to_payloads,
@@ -83,9 +84,11 @@ def get_identity_state(passed) -> tuple[str, bool]:
 
 
 def add_context_tags(properties):
+    properties = properties or {}
     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
@@ -311,6 +314,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.
@@ -321,12 +325,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 {}
 
@@ -337,6 +348,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.
@@ -347,6 +359,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
@@ -354,10 +368,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 {}
 
@@ -368,6 +387,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.
@@ -378,6 +398,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
@@ -385,20 +407,26 @@ 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)
 
     def get_flags_decision(
         self,
         distinct_id: Optional[ID_TYPES] = None,
-        groups: Optional[dict] = {},
+        groups: Optional[dict] = None,
         person_properties=None,
         group_properties=None,
         disable_geoip=None,
+        flag_keys_to_evaluate: Optional[list[str]] = None,
     ) -> FlagsResponse:
         """
         Get feature flags decision.
@@ -409,6 +437,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
@@ -416,8 +446,11 @@ class Client(object):
             ```
 
         Category:
-            Feature Flags
+            Feature flags
         """
+        groups = groups or {}
+        person_properties = person_properties or {}
+        group_properties = group_properties or {}
 
         if distinct_id is None:
             distinct_id = get_context_distinct_id()
@@ -436,6 +469,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,
@@ -505,6 +541,7 @@ class Client(object):
         properties = {**(properties or {}), **system_context()}
 
         properties = add_context_tags(properties)
+        assert properties is not None  # Type hint for mypy
 
         (distinct_id, personless) = get_identity_state(distinct_id)
 
@@ -520,15 +557,37 @@ class Client(object):
         }
 
         if groups:
-            msg["properties"]["$groups"] = groups
+            properties["$groups"] = groups
 
         extra_properties: dict[str, Any] = {}
         feature_variants: Optional[dict[str, Union[bool, str]]] = {}
-        if send_feature_flags:
+
+        # Parse and normalize send_feature_flags parameter
+        flag_options = self._parse_send_feature_flags(send_feature_flags)
+
+        if flag_options["should_send"]:
             try:
-                feature_variants = self.get_feature_variants(
-                    distinct_id, groups, disable_geoip=disable_geoip
-                )
+                if flag_options["only_evaluate_locally"] is True:
+                    # Only use local evaluation
+                    feature_variants = self.get_all_flags(
+                        distinct_id,
+                        groups=(groups or {}),
+                        person_properties=flag_options["person_properties"],
+                        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
+                    feature_variants = self.get_feature_variants(
+                        distinct_id,
+                        groups,
+                        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(
                     f"[FEATURE FLAGS] Unable to get feature variants: {e}"
@@ -555,10 +614,49 @@ class Client(object):
             extra_properties["$active_feature_flags"] = active_feature_flags
 
         if extra_properties:
-            msg["properties"] = {**extra_properties, **msg["properties"]}
+            properties = {**extra_properties, **properties}
+            msg["properties"] = properties
 
         return self._enqueue(msg, disable_geoip)
 
+    def _parse_send_feature_flags(self, send_feature_flags) -> SendFeatureFlagsOptions:
+        """
+        Parse and normalize send_feature_flags parameter into a standard format.
+
+        Args:
+            send_feature_flags: Either bool or SendFeatureFlagsOptions dict
+
+        Returns:
+            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
+        """
+        if isinstance(send_feature_flags, dict):
+            return {
+                "should_send": True,
+                "only_evaluate_locally": send_feature_flags.get(
+                    "only_evaluate_locally"
+                ),
+                "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 {
+                "should_send": send_feature_flags,
+                "only_evaluate_locally": None,
+                "person_properties": None,
+                "group_properties": None,
+                "flag_keys_filter": None,
+            }
+        else:
+            raise TypeError(
+                f"Invalid type for send_feature_flags: {type(send_feature_flags)}. "
+                f"Expected bool or dict."
+            )
+
     def set(self, **kwargs: Unpack[OptionalSetArgs]) -> Optional[str]:
         """
         Set properties on a person profile.
@@ -1071,7 +1169,7 @@ class Client(object):
             ```
 
         Category:
-            Feature Flags
+            Feature flags
         """
         if not self.personal_api_key:
             self.log.warning(
@@ -1097,11 +1195,18 @@ class Client(object):
         feature_flag,
         distinct_id,
         *,
-        groups={},
-        person_properties={},
-        group_properties={},
+        groups=None,
+        person_properties=None,
+        group_properties=None,
         warn_on_unknown_groups=True,
     ) -> FlagValue:
+        groups = groups or {}
+        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")
 
@@ -1117,12 +1222,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"
@@ -1135,11 +1240,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(
@@ -1147,9 +1261,9 @@ class Client(object):
         key,
         distinct_id,
         *,
-        groups={},
-        person_properties={},
-        group_properties={},
+        groups=None,
+        person_properties=None,
+        group_properties=None,
         only_evaluate_locally=False,
         send_feature_flag_events=True,
         disable_geoip=None,
@@ -1177,7 +1291,7 @@ class Client(object):
             ```
 
         Category:
-            Feature Flags
+            Feature flags
         """
         response = self.get_feature_flag(
             key,
@@ -1200,9 +1314,9 @@ class Client(object):
         distinct_id: ID_TYPES,
         *,
         override_match_value: Optional[FlagValue] = None,
-        groups: Dict[str, str] = {},
-        person_properties={},
-        group_properties={},
+        groups: Optional[Dict[str, str]] = None,
+        person_properties=None,
+        group_properties=None,
         only_evaluate_locally=False,
         send_feature_flag_events=True,
         disable_geoip=None,
@@ -1212,9 +1326,16 @@ class Client(object):
 
         person_properties, group_properties = (
             self._add_local_person_and_group_properties(
-                distinct_id, groups, person_properties, group_properties
+                distinct_id,
+                groups or {},
+                person_properties or {},
+                group_properties or {},
             )
         )
+        # Ensure non-None values for type checking
+        groups = groups or {}
+        person_properties = person_properties or {}
+        group_properties = group_properties or {}
 
         flag_result = None
         flag_details = None
@@ -1229,7 +1350,7 @@ class Client(object):
             lookup_match_value = override_match_value or flag_value
             payload = (
                 self._compute_payload_locally(key, lookup_match_value)
-                if lookup_match_value
+                if lookup_match_value is not None
                 else None
             )
             flag_result = FeatureFlagResult.from_value_and_payload(
@@ -1243,7 +1364,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,
@@ -1298,9 +1419,9 @@ class Client(object):
         key,
         distinct_id,
         *,
-        groups={},
-        person_properties={},
-        group_properties={},
+        groups=None,
+        person_properties=None,
+        group_properties=None,
         only_evaluate_locally=False,
         send_feature_flag_events=True,
         disable_geoip=None,
@@ -1348,9 +1469,9 @@ class Client(object):
         key,
         distinct_id,
         *,
-        groups={},
-        person_properties={},
-        group_properties={},
+        groups=None,
+        person_properties=None,
+        group_properties=None,
         only_evaluate_locally=False,
         send_feature_flag_events=True,
         disable_geoip=None,
@@ -1378,7 +1499,7 @@ class Client(object):
             ```
 
         Category:
-            Feature Flags
+            Feature flags
         """
         feature_flag_result = self.get_feature_flag_result(
             key,
@@ -1436,9 +1557,9 @@ class Client(object):
         distinct_id,
         *,
         match_value: Optional[FlagValue] = None,
-        groups={},
-        person_properties={},
-        group_properties={},
+        groups=None,
+        person_properties=None,
+        group_properties=None,
         only_evaluate_locally=False,
         send_feature_flag_events=True,
         disable_geoip=None,
@@ -1468,7 +1589,7 @@ class Client(object):
             ```
 
         Category:
-            Feature Flags
+            Feature flags
         """
         feature_flag_result = self._get_feature_flag_result(
             key,
@@ -1483,7 +1604,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,
@@ -1493,10 +1614,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")
@@ -1530,7 +1656,7 @@ class Client(object):
                 f"$feature/{key}": response,
             }
 
-            if payload:
+            if payload is not None:
                 # if payload is not a string, json serialize it to a string
                 properties["$feature_flag_payload"] = payload
 
@@ -1571,6 +1697,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,
@@ -1606,11 +1733,12 @@ class Client(object):
         self,
         distinct_id,
         *,
-        groups={},
-        person_properties={},
-        group_properties={},
+        groups=None,
+        person_properties=None,
+        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.
@@ -1622,6 +1750,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
@@ -1629,7 +1759,7 @@ class Client(object):
             ```
 
         Category:
-            Feature Flags
+            Feature flags
         """
         response = self.get_all_flags_and_payloads(
             distinct_id,
@@ -1638,6 +1768,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"]
@@ -1646,11 +1777,12 @@ class Client(object):
         self,
         distinct_id,
         *,
-        groups={},
-        person_properties={},
-        group_properties={},
+        groups=None,
+        person_properties=None,
+        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.
@@ -1662,6 +1794,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
@@ -1669,7 +1803,7 @@ class Client(object):
             ```
 
         Category:
-            Feature Flags
+            Feature flags
         """
         if self.disabled:
             return {"featureFlags": None, "featureFlagPayloads": None}
@@ -1680,14 +1814,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,
@@ -1695,6 +1830,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:
@@ -1709,19 +1845,31 @@ class Client(object):
         distinct_id: ID_TYPES,
         *,
         groups: Dict[str, Union[str, int]],
-        person_properties={},
-        group_properties={},
+        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 {}
+
         if self.feature_flags is None and self.personal_api_key:
             self.load_feature_flags()
 
         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,
@@ -1734,23 +1882,23 @@ class Client(object):
                     matched_payload = self._compute_payload_locally(
                         flag["key"], flags[flag["key"]]
                     )
-                    if matched_payload:
+                    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.
@@ -1861,7 +2009,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
@@ -79,7 +232,13 @@ def match_feature_flag_properties(
             # 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 +260,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 +437,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 +461,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 +495,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 +525,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/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()
 
@@ -88,6 +223,56 @@ def mock_anthropic_response_with_cached_tokens():
     )
 
 
+@pytest.fixture
+def mock_anthropic_response_with_tool_calls():
+    return Message(
+        id="msg_456",
+        type="message",
+        role="assistant",
+        content=[
+            {"type": "text", "text": "I'll help you check the weather."},
+            {"type": "text", "text": " Let me look that up."},
+            {
+                "type": "tool_use",
+                "id": "toolu_abc123",
+                "name": "get_weather",
+                "input": {"location": "San Francisco"},
+            },
+        ],
+        model="claude-3-5-sonnet-20241022",
+        usage=Usage(
+            input_tokens=25,
+            output_tokens=15,
+        ),
+        stop_reason="tool_use",
+        stop_sequence=None,
+    )
+
+
+@pytest.fixture
+def mock_anthropic_response_tool_calls_only():
+    return Message(
+        id="msg_789",
+        type="message",
+        role="assistant",
+        content=[
+            {
+                "type": "tool_use",
+                "id": "toolu_def456",
+                "name": "get_weather",
+                "input": {"location": "New York", "unit": "fahrenheit"},
+            }
+        ],
+        model="claude-3-5-sonnet-20241022",
+        usage=Usage(
+            input_tokens=30,
+            output_tokens=12,
+        ),
+        stop_reason="tool_use",
+        stop_sequence=None,
+    )
+
+
 def test_basic_completion(mock_client, mock_anthropic_response):
     with patch(
         "anthropic.resources.Messages.create", return_value=mock_anthropic_response
@@ -112,7 +297,10 @@ def test_basic_completion(mock_client, mock_anthropic_response):
         assert props["$ai_model"] == "claude-3-opus-20240229"
         assert props["$ai_input"] == [{"role": "user", "content": "Hello"}]
         assert props["$ai_output_choices"] == [
-            {"role": "assistant", "content": "Test response"}
+            {
+                "role": "assistant",
+                "content": [{"type": "text", "text": "Test response"}],
+            }
         ]
         assert props["$ai_input_tokens"] == 20
         assert props["$ai_output_tokens"] == 10
@@ -121,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
@@ -260,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
 
@@ -286,7 +402,9 @@ def test_basic_integration(mock_client):
         {"role": "user", "content": "Foo"},
     ]
     assert props["$ai_output_choices"][0]["role"] == "assistant"
-    assert props["$ai_output_choices"][0]["content"] == "Bar"
+    assert props["$ai_output_choices"][0]["content"] == [
+        {"type": "text", "text": "Bar"}
+    ]
     assert props["$ai_input_tokens"] == 18
     assert props["$ai_output_tokens"] == 1
     assert props["$ai_http_status"] == 200
@@ -294,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
 
@@ -326,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"]
 
@@ -425,7 +552,10 @@ def test_cached_tokens(mock_client, mock_anthropic_response_with_cached_tokens):
         assert props["$ai_model"] == "claude-3-opus-20240229"
         assert props["$ai_input"] == [{"role": "user", "content": "Hello"}]
         assert props["$ai_output_choices"] == [
-            {"role": "assistant", "content": "Test response"}
+            {
+                "role": "assistant",
+                "content": [{"type": "text", "text": "Test response"}],
+            }
         ]
         assert props["$ai_input_tokens"] == 20
         assert props["$ai_output_tokens"] == 10
@@ -434,3 +564,473 @@ def test_cached_tokens(mock_client, mock_anthropic_response_with_cached_tokens):
         assert props["$ai_http_status"] == 200
         assert props["foo"] == "bar"
         assert isinstance(props["$ai_latency"], float)
+
+
+def test_tool_definition(mock_client, mock_anthropic_response):
+    with patch(
+        "anthropic.resources.Messages.create",
+        return_value=mock_anthropic_response,
+    ):
+        client = Anthropic(api_key="test-key", posthog_client=mock_client)
+
+        tools = [
+            {
+                "name": "get_weather",
+                "description": "Get the current weather for a specific location",
+                "input_schema": {
+                    "type": "object",
+                    "properties": {
+                        "location": {
+                            "type": "string",
+                            "description": "The city or location name to get weather for",
+                        }
+                    },
+                    "required": ["location"],
+                },
+            }
+        ]
+
+        response = client.messages.create(
+            model="claude-3-5-sonnet-20241022",
+            max_tokens=200,
+            temperature=0.7,
+            tools=tools,
+            messages=[{"role": "user", "content": "hey"}],
+            posthog_distinct_id="test-id",
+            posthog_properties={"foo": "bar"},
+        )
+
+        assert response == mock_anthropic_response
+        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"
+        assert props["$ai_input"] == [{"role": "user", "content": "hey"}]
+        assert props["$ai_output_choices"] == [
+            {
+                "role": "assistant",
+                "content": [{"type": "text", "text": "Test response"}],
+            }
+        ]
+        assert props["$ai_input_tokens"] == 20
+        assert props["$ai_output_tokens"] == 10
+        assert props["$ai_http_status"] == 200
+        assert props["foo"] == "bar"
+        assert isinstance(props["$ai_latency"], float)
+        # Verify that tools are captured in the $ai_tools property
+        assert props["$ai_tools"] == tools
+
+
+def test_tool_calls_in_output_choices(
+    mock_client, mock_anthropic_response_with_tool_calls
+):
+    with patch(
+        "anthropic.resources.Messages.create",
+        return_value=mock_anthropic_response_with_tool_calls,
+    ):
+        client = Anthropic(api_key="test-key", posthog_client=mock_client)
+        response = client.messages.create(
+            model="claude-3-5-sonnet-20241022",
+            max_tokens=200,
+            messages=[
+                {"role": "user", "content": "What's the weather in San Francisco?"}
+            ],
+            tools=[
+                {
+                    "name": "get_weather",
+                    "description": "Get weather",
+                    "input_schema": {
+                        "type": "object",
+                        "properties": {"location": {"type": "string"}},
+                        "required": ["location"],
+                    },
+                }
+            ],
+            posthog_distinct_id="test-id",
+        )
+
+        assert response == mock_anthropic_response_with_tool_calls
+        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"
+        assert props["$ai_output_choices"] == [
+            {
+                "role": "assistant",
+                "content": [
+                    {"type": "text", "text": "I'll help you check the weather."},
+                    {"type": "text", "text": " Let me look that up."},
+                    {
+                        "type": "function",
+                        "id": "toolu_abc123",
+                        "function": {
+                            "name": "get_weather",
+                            "arguments": {"location": "San Francisco"},
+                        },
+                    },
+                ],
+            }
+        ]
+
+        # Check token usage
+        assert props["$ai_input_tokens"] == 25
+        assert props["$ai_output_tokens"] == 15
+        assert props["$ai_http_status"] == 200
+
+
+def test_tool_calls_only_no_content(
+    mock_client, mock_anthropic_response_tool_calls_only
+):
+    with patch(
+        "anthropic.resources.Messages.create",
+        return_value=mock_anthropic_response_tool_calls_only,
+    ):
+        client = Anthropic(api_key="test-key", posthog_client=mock_client)
+        response = client.messages.create(
+            model="claude-3-5-sonnet-20241022",
+            max_tokens=200,
+            messages=[{"role": "user", "content": "Get weather for New York"}],
+            tools=[
+                {
+                    "name": "get_weather",
+                    "description": "Get weather",
+                    "input_schema": {
+                        "type": "object",
+                        "properties": {
+                            "location": {"type": "string"},
+                            "unit": {"type": "string"},
+                        },
+                        "required": ["location"],
+                    },
+                }
+            ],
+            posthog_distinct_id="test-id",
+        )
+
+        assert response == mock_anthropic_response_tool_calls_only
+        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"
+        assert props["$ai_output_choices"] == [
+            {
+                "role": "assistant",
+                "content": [
+                    {
+                        "type": "function",
+                        "id": "toolu_def456",
+                        "function": {
+                            "name": "get_weather",
+                            "arguments": {"location": "New York", "unit": "fahrenheit"},
+                        },
+                    }
+                ],
+            }
+        ]
+
+        # Check token usage
+        assert props["$ai_input_tokens"] == 30
+        assert props["$ai_output_tokens"] == 12
+        assert props["$ai_http_status"] == 200
+
+
+def test_async_tool_calls_in_output_choices(
+    mock_client, mock_anthropic_response_with_tool_calls
+):
+    import asyncio
+
+    async def mock_async_create(**kwargs):
+        return mock_anthropic_response_with_tool_calls
+
+    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():
+            return await async_client.messages.create(
+                model="claude-3-5-sonnet-20241022",
+                max_tokens=200,
+                messages=[
+                    {"role": "user", "content": "What's the weather in San Francisco?"}
+                ],
+                tools=[
+                    {
+                        "name": "get_weather",
+                        "description": "Get weather",
+                        "input_schema": {
+                            "type": "object",
+                            "properties": {"location": {"type": "string"}},
+                            "required": ["location"],
+                        },
+                    }
+                ],
+                posthog_distinct_id="test-id",
+            )
+
+        response = asyncio.run(run_test())
+
+        assert response == mock_anthropic_response_with_tool_calls
+        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"
+        assert props["$ai_output_choices"] == [
+            {
+                "role": "assistant",
+                "content": [
+                    {"type": "text", "text": "I'll help you check the weather."},
+                    {"type": "text", "text": " Let me look that up."},
+                    {
+                        "type": "function",
+                        "id": "toolu_abc123",
+                        "function": {
+                            "name": "get_weather",
+                            "arguments": {"location": "San Francisco"},
+                        },
+                    },
+                ],
+            }
+        ]
+
+        # Check token usage
+        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()
@@ -56,6 +59,91 @@ def mock_google_genai_client():
         yield mock_client_instance
 
 
+@pytest.fixture
+def mock_gemini_response_with_function_calls():
+    mock_response = MagicMock()
+
+    # Mock usage metadata
+    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
+    mock_function_call = MagicMock()
+    mock_function_call.name = "get_current_weather"
+    mock_function_call.args = {"location": "San Francisco"}
+
+    # Mock text part 1
+    mock_text_part1 = MagicMock()
+    mock_text_part1.text = "I'll check the weather for you."
+    # Make hasattr(part, "text") return True
+    type(mock_text_part1).text = mock_text_part1.text
+
+    # Mock text part 2
+    mock_text_part2 = MagicMock()
+    mock_text_part2.text = " Let me look that up."
+    type(mock_text_part2).text = mock_text_part2.text
+
+    # Mock function call part - need to ensure hasattr() works correctly
+    mock_function_part = MagicMock()
+    mock_function_part.function_call = mock_function_call
+    # Make hasattr(part, "function_call") return True
+    type(mock_function_part).function_call = mock_function_part.function_call
+    # Ensure hasattr(part, "text") returns False for the function part
+    del mock_function_part.text
+
+    # Mock content with 2 text parts and 1 function call part
+    mock_content = MagicMock()
+    mock_content.parts = [mock_text_part1, mock_text_part2, mock_function_part]
+
+    # Mock candidate
+    mock_candidate = MagicMock()
+    mock_candidate.content = mock_content
+    mock_response.candidates = [mock_candidate]
+
+    return mock_response
+
+
+@pytest.fixture
+def mock_gemini_response_function_calls_only():
+    mock_response = MagicMock()
+
+    # Mock usage metadata
+    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
+    mock_function_call = MagicMock()
+    mock_function_call.name = "get_current_weather"
+    mock_function_call.args = {"location": "New York", "unit": "fahrenheit"}
+
+    # Mock function call part (no text part) - need to ensure hasattr() works correctly
+    mock_function_part = MagicMock()
+    mock_function_part.function_call = mock_function_call
+    # Make hasattr(part, "function_call") return True
+    type(mock_function_part).function_call = mock_function_part.function_call
+    # Ensure hasattr(part, "text") returns False for the function part
+    del mock_function_part.text
+
+    # Mock content with only function call part
+    mock_content = MagicMock()
+    mock_content.parts = [mock_function_part]
+
+    # Mock candidate
+    mock_candidate = MagicMock()
+    mock_candidate.content = mock_content
+    mock_response.candidates = [mock_candidate]
+
+    return mock_response
+
+
 def test_new_client_basic_generation(
     mock_client, mock_google_genai_client, mock_gemini_response
 ):
@@ -99,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()
@@ -106,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
@@ -145,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
@@ -221,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"]
@@ -318,3 +515,323 @@ def test_new_client_override_defaults(
     assert props["team"] == "ai"  # from defaults
     assert props["feature"] == "chat"  # from call
     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
+
+    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(
+        model="gemini-2.5-flash",
+        contents=["hey"],
+        config=mock_config,
+        posthog_distinct_id="test-id",
+        posthog_properties={"foo": "bar"},
+    )
+
+    assert response == mock_gemini_response
+    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.5-flash"
+    assert props["$ai_input"] == [{"role": "user", "content": "hey"}]
+    assert props["$ai_output_choices"] == [
+        {
+            "role": "assistant",
+            "content": [{"type": "text", "text": "Test response from Gemini"}],
+        }
+    ]
+    assert props["$ai_input_tokens"] == 20
+    assert props["$ai_output_tokens"] == 10
+    assert props["$ai_http_status"] == 200
+    assert props["foo"] == "bar"
+    assert isinstance(props["$ai_latency"], float)
+    # Verify that tools are captured in the $ai_tools property
+    assert props["$ai_tools"] == [mock_tool]
+
+
+def test_function_calls_in_output_choices(
+    mock_client, mock_google_genai_client, mock_gemini_response_with_function_calls
+):
+    """Test that function calls are properly included in $ai_output_choices"""
+    mock_google_genai_client.models.generate_content.return_value = (
+        mock_gemini_response_with_function_calls
+    )
+
+    client = Client(api_key="test-key", posthog_client=mock_client)
+
+    response = client.models.generate_content(
+        model="gemini-2.5-flash",
+        contents=["What's the weather in San Francisco?"],
+        posthog_distinct_id="test-id",
+    )
+
+    assert response == mock_gemini_response_with_function_calls
+    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.5-flash"
+    assert props["$ai_output_choices"] == [
+        {
+            "role": "assistant",
+            "content": [
+                {"type": "text", "text": "I'll check the weather for you."},
+                {"type": "text", "text": " Let me look that up."},
+                {
+                    "type": "function",
+                    "function": {
+                        "name": "get_current_weather",
+                        "arguments": {"location": "San Francisco"},
+                    },
+                },
+            ],
+        }
+    ]
+
+    # Check token usage
+    assert props["$ai_input_tokens"] == 25
+    assert props["$ai_output_tokens"] == 15
+    assert props["$ai_http_status"] == 200
+
+
+def test_function_calls_only_no_content(
+    mock_client, mock_google_genai_client, mock_gemini_response_function_calls_only
+):
+    """Test function calls without text content in $ai_output_choices"""
+    mock_google_genai_client.models.generate_content.return_value = (
+        mock_gemini_response_function_calls_only
+    )
+
+    client = Client(api_key="test-key", posthog_client=mock_client)
+
+    response = client.models.generate_content(
+        model="gemini-2.5-flash",
+        contents=["Get weather for New York"],
+        posthog_distinct_id="test-id",
+    )
+
+    assert response == mock_gemini_response_function_calls_only
+    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.5-flash"
+    assert props["$ai_output_choices"] == [
+        {
+            "role": "assistant",
+            "content": [
+                {
+                    "type": "function",
+                    "function": {
+                        "name": "get_current_weather",
+                        "arguments": {"location": "New York", "unit": "fahrenheit"},
+                    },
+                }
+            ],
+        }
+    ]
+
+    # Check token usage
+    assert props["$ai_input_tokens"] == 30
+    assert props["$ai_output_tokens"] == 12
+    assert props["$ai_http_status"] == 200
+
+
+def test_cache_and_reasoning_tokens(mock_client, mock_google_genai_client):
+    """Test that cache and reasoning tokens are properly extracted"""
+    # Create a mock response with cache and reasoning tokens
+    mock_response = MagicMock()
+    mock_response.text = "Test response with cache"
+
+    mock_usage = MagicMock()
+    mock_usage.prompt_token_count = 100
+    mock_usage.candidates_token_count = 50
+    mock_usage.cached_content_token_count = 30  # Cache tokens
+    mock_usage.thoughts_token_count = 10  # Reasoning tokens
+    mock_response.usage_metadata = mock_usage
+
+    # Mock candidates
+    mock_candidate = MagicMock()
+    mock_candidate.text = "Test response with cache"
+    mock_response.candidates = [mock_candidate]
+
+    mock_google_genai_client.models.generate_content.return_value = mock_response
+
+    client = Client(api_key="test-key", posthog_client=mock_client)
+
+    response = client.models.generate_content(
+        model="gemini-2.5-pro",
+        contents="Test with cache",
+        posthog_distinct_id="test-id",
+    )
+
+    assert response == mock_response
+    assert mock_client.capture.call_count == 1
+
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+
+    # Check that all token types are present
+    assert props["$ai_input_tokens"] == 100
+    assert props["$ai_output_tokens"] == 50
+    assert props["$ai_cache_read_input_tokens"] == 30
+    assert props["$ai_reasoning_tokens"] == 10
+
+
+def test_streaming_cache_and_reasoning_tokens(mock_client, mock_google_genai_client):
+    """Test that cache and reasoning tokens are properly extracted in streaming"""
+    # Create mock chunks with cache and reasoning tokens
+    chunk1 = MagicMock()
+    chunk1.text = "Hello "
+    chunk1_usage = MagicMock()
+    chunk1_usage.prompt_token_count = 100
+    chunk1_usage.candidates_token_count = 5
+    chunk1_usage.cached_content_token_count = 30  # Cache tokens
+    chunk1_usage.thoughts_token_count = 0
+    chunk1.usage_metadata = chunk1_usage
+
+    chunk2 = MagicMock()
+    chunk2.text = "world!"
+    chunk2_usage = MagicMock()
+    chunk2_usage.prompt_token_count = 100
+    chunk2_usage.candidates_token_count = 10
+    chunk2_usage.cached_content_token_count = 30  # Same cache tokens
+    chunk2_usage.thoughts_token_count = 5  # Reasoning tokens
+    chunk2.usage_metadata = chunk2_usage
+
+    mock_stream = iter([chunk1, chunk2])
+    mock_google_genai_client.models.generate_content_stream.return_value = mock_stream
+
+    client = Client(api_key="test-key", posthog_client=mock_client)
+
+    response = client.models.generate_content_stream(
+        model="gemini-2.5-pro",
+        contents="Test streaming with cache",
+        posthog_distinct_id="test-id",
+    )
+
+    # Consume the stream
+    result = list(response)
+    assert len(result) == 2
+
+    # Check PostHog capture was called
+    assert mock_client.capture.call_count == 1
+
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+
+    # Check that all token types are present (should use final chunk's usage)
+    assert props["$ai_input_tokens"] == 100
+    assert props["$ai_output_tokens"] == 10
+    assert props["$ai_cache_read_input_tokens"] == 30
+    assert props["$ai_reasoning_tokens"] == 5

posthog/test/ai/langchain/test_callbacks.py

@@ -5,7 +5,7 @@ import os
 import time
 import uuid
 from typing import List, Literal, Optional, TypedDict, Union
-from unittest.mock import patch
+from unittest.mock import patch, MagicMock
 
 import pytest
 
@@ -1727,3 +1727,152 @@ def test_openai_reasoning_tokens(mock_client):
     assert call["properties"]["$ai_reasoning_tokens"] is not None
     assert call["properties"]["$ai_input_tokens"] is not None
     assert call["properties"]["$ai_output_tokens"] is not None
+
+
+def test_callback_handler_without_client():
+    """Test that CallbackHandler works properly when no PostHog client is passed."""
+    with patch("posthog.ai.langchain.callbacks.setup") as mock_setup:
+        mock_client = mock_setup.return_value
+
+        callbacks = CallbackHandler()
+
+        # Verify that setup() was called
+        mock_setup.assert_called_once()
+
+        # Verify that the client was set to the result of setup()
+        assert callbacks._ph_client == mock_client
+
+        # Test that the callback handler works with a simple chain
+        prompt = ChatPromptTemplate.from_messages([("user", "Foo")])
+        model = FakeMessagesListChatModel(responses=[AIMessage(content="Bar")])
+        chain = prompt | model
+
+        # This should work and call the mock client
+        result = chain.invoke({}, config={"callbacks": [callbacks]})
+        assert result.content == "Bar"
+
+        # Verify that the mock client was used for capturing events
+        assert mock_client.capture.call_count == 3
+
+
+def test_convert_message_to_dict_tool_calls():
+    """Test that _convert_message_to_dict properly converts tool calls in AIMessage."""
+    from posthog.ai.langchain.callbacks import _convert_message_to_dict
+    from langchain_core.messages import AIMessage
+    from langchain_core.messages.tool import ToolCall
+
+    # Create an AIMessage with tool calls
+    tool_calls = [
+        ToolCall(
+            id="call_123",
+            name="get_weather",
+            args={"city": "San Francisco", "units": "celsius"},
+        )
+    ]
+
+    ai_message = AIMessage(
+        content="I'll check the weather for you.", tool_calls=tool_calls
+    )
+
+    # Convert to dict
+    result = _convert_message_to_dict(ai_message)
+
+    # Verify the conversion
+    assert result["role"] == "assistant"
+    assert result["content"] == "I'll check the weather for you."
+    assert result["tool_calls"] == [
+        {
+            "type": "function",
+            "id": "call_123",
+            "function": {
+                "name": "get_weather",
+                "arguments": '{"city": "San Francisco", "units": "celsius"}',
+            },
+        }
+    ]
+
+
+def test_tool_definition(mock_client):
+    """Test that tools defined in invocation parameters are captured in $ai_tools property"""
+    callbacks = CallbackHandler(mock_client)
+    run_id = uuid.uuid4()
+
+    # Define tools to be passed to the invocation parameters
+    tools = [
+        {
+            "type": "function",
+            "function": {
+                "name": "get_weather",
+                "description": "Get the current weather for a specific location",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "location": {
+                            "type": "string",
+                            "description": "The city or location name to get weather for",
+                        }
+                    },
+                    "required": ["location"],
+                },
+            },
+        }
+    ]
+
+    with patch("time.time", return_value=1234567890):
+        callbacks._set_llm_metadata(
+            {"kwargs": {"openai_api_base": "https://api.openai.com/v1"}},
+            run_id,
+            messages=[{"role": "user", "content": "hey"}],
+            invocation_params={"temperature": 0.7, "tools": tools},
+            metadata={"ls_model_name": "gpt-4o-mini", "ls_provider": "openai"},
+            name="test",
+        )
+
+    expected = GenerationMetadata(
+        model="gpt-4o-mini",
+        input=[{"role": "user", "content": "hey"}],
+        start_time=1234567890,
+        model_params={"temperature": 0.7},
+        provider="openai",
+        base_url="https://api.openai.com/v1",
+        name="test",
+        tools=tools,
+        end_time=None,
+    )
+    assert callbacks._runs[run_id] == expected
+
+    with patch("time.time", return_value=1234567891):
+        run = callbacks._pop_run_metadata(run_id)
+    expected.end_time = 1234567891
+    assert run == expected
+    assert callbacks._runs == {}
+
+    # Now test that the tools are properly captured in the PostHog event
+    mock_response = MagicMock()
+    mock_response.generations = [[MagicMock()]]
+
+    callbacks._capture_generation(
+        trace_id=run_id,
+        run_id=run_id,
+        run=run,
+        output=mock_response,
+        parent_run_id=None,
+    )
+
+    assert mock_client.capture.call_count == 1
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+
+    assert call_args["distinct_id"] == run_id
+    assert call_args["event"] == "$ai_generation"
+    assert props["$ai_provider"] == "openai"
+    assert props["$ai_model"] == "gpt-4o-mini"
+    assert props["$ai_input"] == [{"role": "user", "content": "hey"}]
+    assert props["$ai_model_parameters"] == {"temperature": 0.7}
+    assert props["$ai_base_url"] == "https://api.openai.com/v1"
+    assert props["$ai_span_name"] == "test"
+    assert props["$ai_span_id"] == run_id
+    assert props["$ai_trace_id"] == run_id
+    assert props["$ai_latency"] == 1.0
+    # Verify that tools are captured in the $ai_tools property
+    assert props["$ai_tools"] == tools

posthog/test/ai/openai/test_openai.py

@@ -1,4 +1,3 @@
-import json
 import time
 from unittest.mock import patch
 
@@ -26,6 +25,7 @@ try:
         ResponseOutputMessage,
         ResponseOutputText,
         ResponseUsage,
+        ResponseFunctionToolCall,
         ParsedResponse,
     )
     from openai.types.responses.parsed_response import (
@@ -227,11 +227,27 @@ def mock_openai_response_with_tool_calls():
         created=int(time.time()),
         choices=[
             Choice(
-                finish_reason="tool_calls",
+                finish_reason="stop",
                 index=0,
                 message=ChatCompletionMessage(
                     content="I'll check the weather for you.",
                     role="assistant",
+                ),
+            ),
+            Choice(
+                finish_reason="stop",
+                index=1,
+                message=ChatCompletionMessage(
+                    content=" Let me look that up.",
+                    role="assistant",
+                ),
+            ),
+            Choice(
+                finish_reason="tool_calls",
+                index=2,
+                message=ChatCompletionMessage(
+                    content=None,
+                    role="assistant",
                     tool_calls=[
                         ChatCompletionMessageToolCall(
                             id="call_abc123",
@@ -243,7 +259,7 @@ def mock_openai_response_with_tool_calls():
                         )
                     ],
                 ),
-            )
+            ),
         ],
         usage=CompletionUsage(
             completion_tokens=15,
@@ -253,6 +269,97 @@ def mock_openai_response_with_tool_calls():
     )
 
 
+@pytest.fixture
+def mock_openai_response_tool_calls_only():
+    return ChatCompletion(
+        id="test",
+        model="gpt-4",
+        object="chat.completion",
+        created=int(time.time()),
+        choices=[
+            Choice(
+                finish_reason="tool_calls",
+                index=0,
+                message=ChatCompletionMessage(
+                    content=None,
+                    role="assistant",
+                    tool_calls=[
+                        ChatCompletionMessageToolCall(
+                            id="call_def456",
+                            type="function",
+                            function=Function(
+                                name="get_weather",
+                                arguments='{"location": "New York"}',
+                            ),
+                        )
+                    ],
+                ),
+            )
+        ],
+        usage=CompletionUsage(
+            completion_tokens=10,
+            prompt_tokens=25,
+            total_tokens=35,
+        ),
+    )
+
+
+@pytest.fixture
+def mock_responses_api_with_tool_calls():
+    return Response(
+        id="resp_123",
+        object="response",
+        created_at=int(time.time()),
+        model="gpt-4o-mini",
+        status="completed",
+        error=None,
+        incomplete_details=None,
+        instructions=None,
+        max_output_tokens=None,
+        tools=[],
+        tool_choice="auto",
+        parallel_tool_calls=True,
+        output=[
+            ResponseOutputMessage(
+                id="msg_456",
+                type="message",
+                role="assistant",
+                status="completed",
+                content=[
+                    ResponseOutputText(
+                        type="output_text",
+                        text="I'll help you with the weather.",
+                        annotations=[],
+                    ),
+                    ResponseOutputText(
+                        type="output_text",
+                        text=" Let me check that for you.",
+                        annotations=[],
+                    ),
+                ],
+            ),
+            ResponseFunctionToolCall(
+                id="fc_789",
+                type="function_call",
+                name="get_weather",
+                call_id="call_xyz789",
+                arguments='{"location": "Chicago"}',
+                status="completed",
+            ),
+        ],
+        usage=ResponseUsage(
+            input_tokens=30,
+            output_tokens=20,
+            input_tokens_details={"prompt_tokens": 30, "cached_tokens": 0},
+            output_tokens_details={"reasoning_tokens": 0},
+            total_tokens=50,
+        ),
+        previous_response_id=None,
+        user=None,
+        metadata={},
+    )
+
+
 def test_basic_completion(mock_client, mock_openai_response):
     with patch(
         "openai.resources.chat.completions.Completions.create",
@@ -278,7 +385,10 @@ def test_basic_completion(mock_client, mock_openai_response):
         assert props["$ai_model"] == "gpt-4"
         assert props["$ai_input"] == [{"role": "user", "content": "Hello"}]
         assert props["$ai_output_choices"] == [
-            {"role": "assistant", "content": "Test response"}
+            {
+                "role": "assistant",
+                "content": [{"type": "text", "text": "Test response"}],
+            }
         ]
         assert props["$ai_input_tokens"] == 20
         assert props["$ai_output_tokens"] == 10
@@ -427,7 +537,10 @@ def test_cached_tokens(mock_client, mock_openai_response_with_cached_tokens):
         assert props["$ai_model"] == "gpt-4"
         assert props["$ai_input"] == [{"role": "user", "content": "Hello"}]
         assert props["$ai_output_choices"] == [
-            {"role": "assistant", "content": "Test response"}
+            {
+                "role": "assistant",
+                "content": [{"type": "text", "text": "Test response"}],
+            }
         ]
         assert props["$ai_input_tokens"] == 20
         assert props["$ai_output_tokens"] == 10
@@ -475,24 +588,34 @@ def test_tool_calls(mock_client, mock_openai_response_with_tool_calls):
             {"role": "user", "content": "What's the weather in San Francisco?"}
         ]
         assert props["$ai_output_choices"] == [
-            {"role": "assistant", "content": "I'll check the weather for you."}
+            {
+                "role": "assistant",
+                "content": [
+                    {"type": "text", "text": "I'll check the weather for you."},
+                    {"type": "text", "text": " Let me look that up."},
+                    {
+                        "type": "function",
+                        "id": "call_abc123",
+                        "function": {
+                            "name": "get_weather",
+                            "arguments": '{"location": "San Francisco", "unit": "celsius"}',
+                        },
+                    },
+                ],
+            }
         ]
 
-        # Check that tool calls are properly captured
+        # Check that defined tools are properly captured in $ai_tools
         assert "$ai_tools" in props
-        tool_calls = props["$ai_tools"]
-        assert len(tool_calls) == 1
+        defined_tools = props["$ai_tools"]
+        assert len(defined_tools) == 1
 
-        # Verify the tool call details
-        tool_call = tool_calls[0]
-        assert tool_call.id == "call_abc123"
-        assert tool_call.type == "function"
-        assert tool_call.function.name == "get_weather"
-
-        # Verify the arguments
-        arguments = tool_call.function.arguments
-        parsed_args = json.loads(arguments)
-        assert parsed_args == {"location": "San Francisco", "unit": "celsius"}
+        # Verify the defined tool details
+        defined_tool = defined_tools[0]
+        assert defined_tool["type"] == "function"
+        assert defined_tool["function"]["name"] == "get_weather"
+        assert defined_tool["function"]["description"] == "Get weather"
+        assert defined_tool["function"]["parameters"] == {}
 
         # Check token usage
         assert props["$ai_input_tokens"] == 20
@@ -500,6 +623,117 @@ def test_tool_calls(mock_client, mock_openai_response_with_tool_calls):
         assert props["$ai_http_status"] == 200
 
 
+def test_tool_calls_only_no_content(mock_client, mock_openai_response_tool_calls_only):
+    with patch(
+        "openai.resources.chat.completions.Completions.create",
+        return_value=mock_openai_response_tool_calls_only,
+    ):
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+        response = client.chat.completions.create(
+            model="gpt-4",
+            messages=[{"role": "user", "content": "Get weather for New York"}],
+            tools=[
+                {
+                    "type": "function",
+                    "function": {
+                        "name": "get_weather",
+                        "description": "Get weather",
+                        "parameters": {},
+                    },
+                }
+            ],
+            posthog_distinct_id="test-id",
+        )
+
+        assert response == mock_openai_response_tool_calls_only
+        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_choices"] == [
+            {
+                "role": "assistant",
+                "content": [
+                    {
+                        "type": "function",
+                        "id": "call_def456",
+                        "function": {
+                            "name": "get_weather",
+                            "arguments": '{"location": "New York"}',
+                        },
+                    }
+                ],
+            }
+        ]
+
+        # Check token usage
+        assert props["$ai_input_tokens"] == 25
+        assert props["$ai_output_tokens"] == 10
+        assert props["$ai_http_status"] == 200
+
+
+def test_responses_api_tool_calls(mock_client, mock_responses_api_with_tool_calls):
+    with patch(
+        "openai.resources.responses.Responses.create",
+        return_value=mock_responses_api_with_tool_calls,
+    ):
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+        response = client.responses.create(
+            model="gpt-4o-mini",
+            input=[{"role": "user", "content": "What's the weather in Chicago?"}],
+            tools=[
+                {
+                    "name": "get_weather",
+                    "type": "function",
+                    "function": {
+                        "name": "get_weather",
+                        "description": "Get weather",
+                        "parameters": {},
+                    },
+                }
+            ],
+            posthog_distinct_id="test-id",
+        )
+
+        assert response == mock_responses_api_with_tool_calls
+        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_output_choices"] == [
+            {
+                "role": "assistant",
+                "content": [
+                    {"type": "text", "text": "I'll help you with the weather."},
+                    {"type": "text", "text": " Let me check that for you."},
+                    {
+                        "type": "function",
+                        "id": "call_xyz789",
+                        "function": {
+                            "name": "get_weather",
+                            "arguments": '{"location": "Chicago"}',
+                        },
+                    },
+                ],
+            }
+        ]
+
+        # Check token usage
+        assert props["$ai_input_tokens"] == 30
+        assert props["$ai_output_tokens"] == 20
+        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 = [
@@ -644,26 +878,41 @@ def test_streaming_with_tool_calls(mock_client):
         assert props["$ai_provider"] == "openai"
         assert props["$ai_model"] == "gpt-4"
 
-        # Check that the tool calls were properly accumulated
+        # Check that defined tools are properly captured in $ai_tools
         assert "$ai_tools" in props
-        tool_calls = props["$ai_tools"]
-        assert len(tool_calls) == 1
+        defined_tools = props["$ai_tools"]
+        assert len(defined_tools) == 1
 
-        # Verify the complete tool call was properly assembled
-        tool_call = tool_calls[0]
-        assert tool_call.id == "call_abc123"
-        assert tool_call.type == "function"
-        assert tool_call.function.name == "get_weather"
+        # Verify the defined tool details
+        defined_tool = defined_tools[0]
+        assert defined_tool["type"] == "function"
+        assert defined_tool["function"]["name"] == "get_weather"
+        assert defined_tool["function"]["description"] == "Get weather"
+        assert defined_tool["function"]["parameters"] == {}
 
-        # Verify the arguments were concatenated correctly
-        arguments = tool_call.function.arguments
-        parsed_args = json.loads(arguments)
-        assert parsed_args == {"location": "San Francisco", "unit": "celsius"}
+        # Check that both text content and tool calls were accumulated
+        output_content = props["$ai_output_choices"][0]["content"]
 
-        # Check that the content was also accumulated
+        # 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
@@ -696,7 +945,10 @@ def test_responses_api(mock_client, mock_openai_response_with_responses_api):
         assert props["$ai_model"] == "gpt-4o-mini"
         assert props["$ai_input"] == [{"role": "user", "content": "Hello"}]
         assert props["$ai_output_choices"] == [
-            {"role": "assistant", "content": "Test response"}
+            {
+                "role": "assistant",
+                "content": [{"type": "text", "text": "Test response"}],
+            }
         ]
         assert props["$ai_input_tokens"] == 10
         assert props["$ai_output_tokens"] == 10
@@ -765,7 +1017,12 @@ def test_responses_parse(mock_client, mock_parsed_response):
         assert props["$ai_output_choices"] == [
             {
                 "role": "assistant",
-                "content": '{"name": "Science Fair", "date": "Friday", "participants": ["Alice", "Bob"]}',
+                "content": [
+                    {
+                        "type": "text",
+                        "text": '{"name": "Science Fair", "date": "Friday", "participants": ["Alice", "Bob"]}',
+                    }
+                ],
             }
         ]
         assert props["$ai_input_tokens"] == 15
@@ -774,3 +1031,145 @@ def test_responses_parse(mock_client, mock_parsed_response):
         assert props["$ai_http_status"] == 200
         assert props["foo"] == "bar"
         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)
+
+
+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(
+        "openai.resources.chat.completions.Completions.create",
+        return_value=mock_openai_response,
+    ):
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+
+        # Define tools to be passed to the create function
+        tools = [
+            {
+                "type": "function",
+                "function": {
+                    "name": "get_weather",
+                    "description": "Get the current weather for a specific location",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "location": {
+                                "type": "string",
+                                "description": "The city or location name to get weather for",
+                            }
+                        },
+                        "required": ["location"],
+                    },
+                },
+            }
+        ]
+
+        response = client.chat.completions.create(
+            model="gpt-4o-mini",
+            messages=[{"role": "user", "content": "hey"}],
+            tools=tools,
+            posthog_distinct_id="test-id",
+            posthog_properties={"foo": "bar"},
+        )
+
+        assert response == mock_openai_response
+        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"] == [{"role": "user", "content": "hey"}]
+        assert props["$ai_output_choices"] == [
+            {
+                "role": "assistant",
+                "content": [{"type": "text", "text": "Test response"}],
+            }
+        ]
+        assert props["$ai_input_tokens"] == 20
+        assert props["$ai_output_tokens"] == 10
+        assert props["$ai_http_status"] == 200
+        assert props["foo"] == "bar"
+        assert isinstance(props["$ai_latency"], float)
+        # Verify that tools are captured in the $ai_tools property
+        assert props["$ai_tools"] == tools

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/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):
@@ -647,8 +648,8 @@ class TestClient(unittest.TestCase):
                 timeout=3,
                 distinct_id="distinct_id",
                 groups={},
-                person_properties=None,
-                group_properties=None,
+                person_properties={},
+                group_properties={},
                 geoip_disable=True,
             )
 
@@ -711,8 +712,8 @@ class TestClient(unittest.TestCase):
                 timeout=12,
                 distinct_id="distinct_id",
                 groups={},
-                person_properties=None,
-                group_properties=None,
+                person_properties={},
+                group_properties={},
                 geoip_disable=False,
             )
 
@@ -751,6 +752,186 @@ class TestClient(unittest.TestCase):
 
             self.assertEqual(patch_flags.call_count, 0)
 
+    @mock.patch("posthog.client.flags")
+    def test_capture_with_send_feature_flags_options_only_evaluate_locally_true(
+        self, patch_flags
+    ):
+        """Test that SendFeatureFlagsOptions with only_evaluate_locally=True uses local evaluation"""
+        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,
+            )
+
+            # Set up local flags
+            client.feature_flags = [
+                {
+                    "id": 1,
+                    "key": "local-flag",
+                    "active": True,
+                    "filters": {
+                        "groups": [
+                            {
+                                "properties": [{"key": "region", "value": "US"}],
+                                "rollout_percentage": 100,
+                            }
+                        ],
+                    },
+                }
+            ]
+
+            send_options = {
+                "only_evaluate_locally": True,
+                "person_properties": {"region": "US"},
+            }
+
+            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 not called (no remote evaluation)
+            patch_flags.assert_not_called()
+
+            # Check the message includes the local flag
+            mock_post.assert_called_once()
+            batch_data = mock_post.call_args[1]["batch"]
+            msg = batch_data[0]
+
+            self.assertEqual(msg["properties"]["$feature/local-flag"], True)
+            self.assertEqual(msg["properties"]["$active_feature_flags"], ["local-flag"])
+
+    @mock.patch("posthog.client.flags")
+    def test_capture_with_send_feature_flags_options_only_evaluate_locally_false(
+        self, patch_flags
+    ):
+        """Test that SendFeatureFlagsOptions with only_evaluate_locally=False forces remote evaluation"""
+        patch_flags.return_value = {"featureFlags": {"remote-flag": "remote-value"}}
+
+        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 = {
+                "only_evaluate_locally": False,
+                "person_properties": {"plan": "premium"},
+                "group_properties": {"company": {"type": "enterprise"}},
+            }
+
+            msg_uuid = client.capture(
+                "test event",
+                distinct_id="distinct_id",
+                groups={"company": "acme"},
+                send_feature_flags=send_options,
+            )
+
+            self.assertIsNotNone(msg_uuid)
+            self.assertFalse(self.failed)
+
+            # Verify flags() was called with the correct properties
+            patch_flags.assert_called_once()
+            call_args = patch_flags.call_args[1]
+            self.assertEqual(call_args["person_properties"], {"plan": "premium"})
+            self.assertEqual(
+                call_args["group_properties"], {"company": {"type": "enterprise"}}
+            )
+
+            # Check the message includes the remote flag
+            mock_post.assert_called_once()
+            batch_data = mock_post.call_args[1]["batch"]
+            msg = batch_data[0]
+
+            self.assertEqual(msg["properties"]["$feature/remote-flag"], "remote-value")
+
+    @mock.patch("posthog.client.flags")
+    def test_capture_with_send_feature_flags_options_default_behavior(
+        self, patch_flags
+    ):
+        """Test that SendFeatureFlagsOptions without only_evaluate_locally defaults to remote evaluation"""
+        patch_flags.return_value = {"featureFlags": {"default-flag": "default-value"}}
+
+        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 = {
+                "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 (default to remote evaluation)
+            patch_flags.assert_called_once()
+            call_args = patch_flags.call_args[1]
+            self.assertEqual(call_args["person_properties"], {"subscription": "pro"})
+
+            # Check the message includes the flag
+            mock_post.assert_called_once()
+            batch_data = mock_post.call_args[1]["batch"]
+            msg = batch_data[0]
+
+            self.assertEqual(
+                msg["properties"]["$feature/default-flag"], "default-value"
+            )
+
+    @mock.patch("posthog.client.flags")
+    def test_capture_exception_with_send_feature_flags_options(self, patch_flags):
+        """Test that capture_exception also supports SendFeatureFlagsOptions"""
+        patch_flags.return_value = {"featureFlags": {"exception-flag": True}}
+
+        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 = {
+                "only_evaluate_locally": False,
+                "person_properties": {"user_type": "admin"},
+            }
+
+            try:
+                raise ValueError("Test exception")
+            except ValueError as e:
+                msg_uuid = client.capture_exception(
+                    e, distinct_id="distinct_id", send_feature_flags=send_options
+                )
+
+            self.assertIsNotNone(msg_uuid)
+            self.assertFalse(self.failed)
+
+            # Verify flags() was called with the correct properties
+            patch_flags.assert_called_once()
+            call_args = patch_flags.call_args[1]
+            self.assertEqual(call_args["person_properties"], {"user_type": "admin"})
+
+            # Check the message includes the flag
+            mock_post.assert_called_once()
+            batch_data = mock_post.call_args[1]["batch"]
+            msg = batch_data[0]
+
+            self.assertEqual(msg["event"], "$exception")
+            self.assertEqual(msg["properties"]["$feature/exception-flag"], True)
+
     def test_stringifies_distinct_id(self):
         # A large number that loses precision in node:
         # node -e "console.log(157963456373623802 + 1)" > 157963456373623800
@@ -1561,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(
@@ -1575,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")
@@ -1591,7 +1774,7 @@ class TestClient(unittest.TestCase):
 
     @mock.patch("posthog.client.Poller")
     @mock.patch("posthog.client.get")
-    def test_call_identify_fails(self, patch_get, patch_poll):
+    def test_call_identify_fails(self, patch_get, patch_poller):
         def raise_effect():
             raise Exception("http exception")
 
@@ -1635,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()
@@ -1661,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()
@@ -1877,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)
@@ -1978,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,
@@ -1993,3 +2179,247 @@ class TestClient(unittest.TestCase):
         result = client.get_remote_config_payload("test-flag")
 
         self.assertIsNone(result)
+
+    def test_parse_send_feature_flags_method(self):
+        """Test the _parse_send_feature_flags helper method"""
+        client = Client(FAKE_TEST_API_KEY, sync_mode=True)
+
+        # Test boolean True
+        result = client._parse_send_feature_flags(True)
+        expected = {
+            "should_send": True,
+            "only_evaluate_locally": None,
+            "person_properties": None,
+            "group_properties": None,
+            "flag_keys_filter": None,
+        }
+        self.assertEqual(result, expected)
+
+        # Test boolean False
+        result = client._parse_send_feature_flags(False)
+        expected = {
+            "should_send": False,
+            "only_evaluate_locally": None,
+            "person_properties": None,
+            "group_properties": None,
+            "flag_keys_filter": None,
+        }
+        self.assertEqual(result, expected)
+
+        # Test options dict with all fields
+        options = {
+            "only_evaluate_locally": True,
+            "person_properties": {"plan": "premium"},
+            "group_properties": {"company": {"type": "enterprise"}},
+        }
+        result = client._parse_send_feature_flags(options)
+        expected = {
+            "should_send": True,
+            "only_evaluate_locally": True,
+            "person_properties": {"plan": "premium"},
+            "group_properties": {"company": {"type": "enterprise"}},
+            "flag_keys_filter": None,
+        }
+        self.assertEqual(result, expected)
+
+        # Test options dict with partial fields
+        options = {"person_properties": {"user_id": "123"}}
+        result = client._parse_send_feature_flags(options)
+        expected = {
+            "should_send": True,
+            "only_evaluate_locally": None,
+            "person_properties": {"user_id": "123"},
+            "group_properties": None,
+            "flag_keys_filter": None,
+        }
+        self.assertEqual(result, expected)
+
+        # Test empty dict
+        result = client._parse_send_feature_flags({})
+        expected = {
+            "should_send": True,
+            "only_evaluate_locally": None,
+            "person_properties": None,
+            "group_properties": None,
+            "flag_keys_filter": None,
+        }
+        self.assertEqual(result, expected)
+
+        # Test invalid types
+        with self.assertRaises(TypeError) as cm:
+            client._parse_send_feature_flags("invalid")
+        self.assertIn("Invalid type for send_feature_flags", str(cm.exception))
+
+        with self.assertRaises(TypeError) as cm:
+            client._parse_send_feature_flags(123)
+        self.assertIn("Invalid type for send_feature_flags", str(cm.exception))
+
+        with self.assertRaises(TypeError) as cm:
+            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"""
+        client = Client(
+            FAKE_TEST_API_KEY,
+            personal_api_key="test_personal_api_key",
+            sync_mode=True,
+        )
+
+        # Set up local evaluation with a flag that has empty string payload
+        client.feature_flags = [
+            {
+                "id": 1,
+                "name": "Test flag",
+                "key": "test-flag",
+                "is_simple_flag": False,
+                "active": True,
+                "rollout_percentage": None,
+                "filters": {
+                    "groups": [
+                        {
+                            "properties": [],
+                            "rollout_percentage": None,
+                            "variant": "empty-variant",
+                        }
+                    ],
+                    "multivariate": {
+                        "variants": [
+                            {
+                                "key": "empty-variant",
+                                "name": "Empty Variant",
+                                "rollout_percentage": 100,
+                            }
+                        ]
+                    },
+                    "payloads": {"empty-variant": ""},  # Empty string payload
+                },
+            }
+        ]
+
+        # Test get_feature_flag_result
+        result = client.get_feature_flag_result(
+            "test-flag", "test-user", only_evaluate_locally=True
+        )
+
+        # Should return a FeatureFlagResult, not None
+        self.assertIsNotNone(result)
+        self.assertEqual(result.key, "test-flag")
+        self.assertEqual(result.get_value(), "empty-variant")
+        self.assertEqual(result.payload, "")  # Should be empty string, not None
+
+    @mock.patch("posthog.client.batch_post")
+    def test_get_all_flags_and_payloads_with_empty_string(self, patch_batch_post):
+        """Test that get_all_flags_and_payloads includes flags with empty string payloads"""
+        client = Client(
+            FAKE_TEST_API_KEY,
+            personal_api_key="test_personal_api_key",
+            sync_mode=True,
+        )
+
+        # Set up multiple flags with different payload types
+        client.feature_flags = [
+            {
+                "id": 1,
+                "name": "Flag with empty payload",
+                "key": "empty-payload-flag",
+                "is_simple_flag": False,
+                "active": True,
+                "filters": {
+                    "groups": [{"properties": [], "variant": "variant1"}],
+                    "multivariate": {
+                        "variants": [{"key": "variant1", "rollout_percentage": 100}]
+                    },
+                    "payloads": {"variant1": ""},  # Empty string
+                },
+            },
+            {
+                "id": 2,
+                "name": "Flag with normal payload",
+                "key": "normal-payload-flag",
+                "is_simple_flag": False,
+                "active": True,
+                "filters": {
+                    "groups": [{"properties": [], "variant": "variant2"}],
+                    "multivariate": {
+                        "variants": [{"key": "variant2", "rollout_percentage": 100}]
+                    },
+                    "payloads": {"variant2": "normal payload"},
+                },
+            },
+        ]
+
+        result = client.get_all_flags_and_payloads(
+            "test-user", only_evaluate_locally=True
+        )
+
+        # Check that both flags are included
+        self.assertEqual(result["featureFlags"]["empty-payload-flag"], "variant1")
+        self.assertEqual(result["featureFlags"]["normal-payload-flag"], "variant2")
+
+        # Check that empty string payload is included (not filtered out)
+        self.assertIn("empty-payload-flag", result["featureFlagPayloads"])
+        self.assertEqual(result["featureFlagPayloads"]["empty-payload-flag"], "")
+        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"])

posthog/types.py

@@ -9,6 +9,27 @@ 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.
+
+    Args:
+        only_evaluate_locally: Whether to only use local evaluation for feature flags.
+            If True, only flags that can be evaluated locally will be included.
+            If False, remote evaluation via /flags API will be used when needed.
+        person_properties: Properties to use for feature flag evaluation specific to this event.
+            These properties will be merged with any existing person properties.
+        group_properties: Group properties to use for feature flag evaluation specific to this event.
+            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)
 class FlagReason:
     code: str
@@ -92,7 +113,7 @@ class FeatureFlag:
             variant=variant,
             reason=None,
             metadata=LegacyFlagMetadata(
-                payload=payload if payload else None,
+                payload=payload,
             ),
         )
 
@@ -160,7 +181,9 @@ class FeatureFlagResult:
             key=key,
             enabled=enabled,
             variant=variant,
-            payload=json.loads(payload) if isinstance(payload, str) else payload,
+            payload=json.loads(payload)
+            if isinstance(payload, str) and payload
+            else payload,
             reason=None,
         )
 
@@ -201,6 +224,7 @@ class FeatureFlagResult:
             payload=(
                 json.loads(details.metadata.payload)
                 if isinstance(details.metadata.payload, str)
+                and details.metadata.payload
                 else details.metadata.payload
             ),
             reason=details.reason.description if details.reason else None,
@@ -278,5 +302,7 @@ def to_payloads(response: FlagsResponse) -> Optional[dict[str, str]]:
     return {
         key: value.metadata.payload
         for key, value in response.get("flags", {}).items()
-        if isinstance(value, FeatureFlag) and value.enabled and value.metadata.payload
+        if isinstance(value, FeatureFlag)
+        and value.enabled
+        and value.metadata.payload is not None
     }

posthog/version.py

@@ -1,4 +1,4 @@
-VERSION = "6.2.1"
+VERSION = "6.7.3"
 
 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()