WIP

posthog/__init__.py

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

posthog/args.py

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

posthog/client.py

@@ -1,16 +1,25 @@
 import atexit
+import json
 import logging
 import os
 import sys
 from datetime import datetime, timedelta
-from typing import Any, Callable, Dict, Optional, Union
+from io import BytesIO
+from typing import Any, Dict, List, Optional, Tuple, Union
 from typing_extensions import Unpack
 from uuid import uuid4
 
 from dateutil.tz import tzutc
 from six import string_types
 
-from posthog.args import OptionalCaptureArgs, OptionalSetArgs, ID_TYPES, ExceptionArg
+from posthog.args import (
+    OptionalCaptureArgs,
+    OptionalSetArgs,
+    OptionalCaptureAIArgs,
+    AI_EVENT_TYPE,
+    ID_TYPES,
+    ExceptionArg,
+)
 from posthog.consumer import Consumer
 from posthog.exception_capture import ExceptionCapture
 from posthog.exception_utils import (
@@ -103,6 +112,106 @@ def add_context_tags(properties):
     return properties
 
 
+def _generate_multipart_boundary() -> str:
+    """Generate a random boundary string for multipart requests."""
+    return f"----WebKitFormBoundary{uuid4().hex[:16]}"
+
+
+def _encode_multipart_part(
+    boundary: str,
+    name: str,
+    content: bytes,
+    content_type: str,
+    filename: Optional[str] = None,
+) -> bytes:
+    """Encode a single part of a multipart/form-data request."""
+    part = f"--{boundary}\r\n".encode("utf-8")
+
+    if filename:
+        part += f'Content-Disposition: form-data; name="{name}"; filename="{filename}"\r\n'.encode(
+            "utf-8"
+        )
+    else:
+        part += f'Content-Disposition: form-data; name="{name}"\r\n'.encode("utf-8")
+
+    part += f"Content-Type: {content_type}\r\n\r\n".encode("utf-8")
+    part += content
+    part += b"\r\n"
+
+    return part
+
+
+def _build_multipart_body(
+    event_data: Dict[str, Any],
+    properties: Optional[Dict[str, Any]],
+    blob_properties: Optional[List[str]],
+) -> Tuple[bytes, str]:
+    """
+    Build a multipart/form-data body for AI capture endpoint.
+
+    Args:
+        event_data: The event data (uuid, event, distinct_id, timestamp)
+        properties: Event properties (small properties)
+        blob_properties: List of property names to send as blobs
+
+    Returns:
+        Tuple of (body_bytes, content_type_header)
+    """
+    boundary = _generate_multipart_boundary()
+    body = BytesIO()
+
+    # Part 1: Event (required, must be first)
+    event_json = json.dumps(event_data).encode("utf-8")
+    body.write(
+        _encode_multipart_part(boundary, "event", event_json, "application/json")
+    )
+
+    # Separate properties into small properties and blob properties
+    small_properties = {}
+    blob_data = {}
+
+    if properties:
+        blob_property_names = set(blob_properties or [])
+        for key, value in properties.items():
+            if key in blob_property_names:
+                blob_data[key] = value
+            else:
+                small_properties[key] = value
+
+    # Part 2: Event properties (if there are any small properties)
+    if small_properties:
+        properties_json = json.dumps(small_properties).encode("utf-8")
+        body.write(
+            _encode_multipart_part(
+                boundary, "event.properties", properties_json, "application/json"
+            )
+        )
+
+    # Part 3+: Blob parts
+    for property_name, property_value in blob_data.items():
+        # Serialize the blob data as JSON
+        blob_json = json.dumps(property_value).encode("utf-8")
+        blob_filename = f"blob_{uuid4().hex[:8]}"
+
+        body.write(
+            _encode_multipart_part(
+                boundary,
+                f"event.properties.{property_name}",
+                blob_json,
+                "application/json",
+                filename=blob_filename,
+            )
+        )
+
+    # Final boundary
+    body.write(f"--{boundary}--\r\n".encode("utf-8"))
+
+    body_bytes = body.getvalue()
+    content_type = f"multipart/form-data; boundary={boundary}"
+
+    return body_bytes, content_type
+
+
 def no_throw(default_return=None):
     """
     Decorator to prevent raising exceptions from public API methods.
@@ -1005,6 +1114,196 @@ class Client(object):
         except Exception as e:
             self.log.exception(f"Failed to capture exception: {e}")
 
+    @no_throw()
+    def capture_ai(
+        self, event: AI_EVENT_TYPE, **kwargs: Unpack[OptionalCaptureAIArgs]
+    ) -> Optional[str]:
+        """
+        Capture an AI event to the dedicated AI endpoint with support for large payloads.
+
+        This method sends AI events (like $ai_generation, $ai_trace, etc.) to PostHog's
+        specialized AI endpoint (/i/v0/ai) which supports large payloads through multipart/form-data
+        and blob storage in S3.
+
+        Args:
+            event: The AI event type. Must be one of: "$ai_generation", "$ai_trace", "$ai_span",
+                   "$ai_embedding", "$ai_metric", "$ai_feedback"
+            distinct_id: The distinct ID of the user.
+            properties: A dictionary of AI event properties. Must include required properties based on event type:
+                - All events: "$ai_model" (required)
+                - $ai_generation: "$ai_provider", "$ai_trace_id" (required)
+                - $ai_trace: "$ai_trace_id" (required)
+                - $ai_span: "$ai_trace_id", "$ai_span_id" (required)
+                - $ai_embedding: "$ai_provider", "$ai_trace_id" (required)
+            blob_properties: List of property names to send as blobs (large data stored in S3).
+                Common blob properties: "$ai_input", "$ai_output_choices", "$ai_input_state", "$ai_output_state"
+                If not provided, defaults to common blob properties based on event type.
+            timestamp: The timestamp of the event.
+            uuid: A unique identifier for the event.
+            groups: A dictionary of group information.
+            disable_geoip: Whether to disable GeoIP for this event.
+
+        Examples:
+            ```python
+            # $ai_generation event with blobs
+            posthog.capture_ai(
+                "$ai_generation",
+                distinct_id="user_123",
+                properties={
+                    "$ai_model": "gpt-4",
+                    "$ai_provider": "openai",
+                    "$ai_trace_id": "trace_abc123",
+                    "$ai_input": {
+                        "messages": [
+                            {"role": "user", "content": "Hello!"}
+                        ]
+                    },
+                    "$ai_output_choices": {
+                        "choices": [{"message": {"role": "assistant", "content": "Hi there!"}}]
+                    },
+                    "$ai_completion_tokens": 150,
+                    "$ai_prompt_tokens": 50
+                },
+                blob_properties=["$ai_input", "$ai_output_choices"]
+            )
+            ```
+
+            ```python
+            # $ai_trace event
+            posthog.capture_ai(
+                "$ai_trace",
+                distinct_id="user_123",
+                properties={
+                    "$ai_model": "gpt-4",
+                    "$ai_trace_id": "trace_abc123"
+                }
+            )
+            ```
+
+            ```python
+            # $ai_metric event
+            posthog.capture_ai(
+                "$ai_metric",
+                distinct_id="user_123",
+                properties={
+                    "$ai_model": "gpt-4",
+                    "$ai_trace_id": "trace_abc123",
+                    "$ai_metric_name": "accuracy",
+                    "$ai_metric_value": "0.95"
+                }
+            )
+            ```
+
+        Category:
+            AI Events
+
+        Note: This method sends events synchronously to the AI endpoint, bypassing the queue system.
+        """
+        import requests
+
+        if self.disabled:
+            return None
+
+        distinct_id = kwargs.get("distinct_id", None)
+        properties = kwargs.get("properties", None)
+        timestamp = kwargs.get("timestamp", None)
+        uuid = kwargs.get("uuid", None)
+        groups = kwargs.get("groups", None)
+        blob_properties = kwargs.get("blob_properties", None)
+
+        # Default blob properties based on event type if not provided
+        if blob_properties is None:
+            default_blob_properties = {
+                "$ai_generation": ["$ai_input", "$ai_output_choices"],
+                "$ai_trace": ["$ai_input_state", "$ai_output_state"],
+                "$ai_span": ["$ai_input_state", "$ai_output_state"],
+                "$ai_embedding": ["$ai_input"],
+            }
+            blob_properties = default_blob_properties.get(event, [])
+
+        properties = properties or {}
+
+        # Get distinct_id
+        (distinct_id, personless) = get_identity_state(distinct_id)
+
+        if personless and "$process_person_profile" not in properties:
+            properties["$process_person_profile"] = False
+
+        # Prepare timestamp
+        if timestamp is None:
+            timestamp = datetime.now(tz=tzutc())
+
+        timestamp = guess_timezone(timestamp)
+        timestamp_str = timestamp.isoformat()
+
+        # Generate UUID if not provided
+        if uuid:
+            uuid_str = stringify_id(uuid)
+        else:
+            uuid_str = stringify_id(uuid4())
+
+        # Add groups to properties
+        if groups:
+            properties["$groups"] = groups
+
+        # Prepare event data (the main event part)
+        event_data = {
+            "uuid": uuid_str,
+            "event": event,
+            "distinct_id": stringify_id(distinct_id),
+            "timestamp": timestamp_str,
+        }
+
+        # Build multipart body
+        body_bytes, content_type = _build_multipart_body(
+            event_data, properties, blob_properties
+        )
+
+        # Send request to AI endpoint
+        url = remove_trailing_slash(self.host) + "/i/v0/ai"
+
+        headers = {
+            "Content-Type": content_type,
+            "Authorization": f"Bearer {self.api_key}",
+            "User-Agent": f"posthog-python/{VERSION}",
+        }
+
+        self.log.debug(f"Sending AI event to {url}: {event} (uuid: {uuid_str})")
+
+        try:
+            response = requests.post(
+                url,
+                data=body_bytes,
+                headers=headers,
+                timeout=self.timeout,
+            )
+
+            if response.status_code == 200:
+                self.log.debug(f"AI event captured successfully: {event}")
+                return uuid_str
+            else:
+                error_message = f"AI capture failed with status {response.status_code}"
+                try:
+                    error_detail = response.json()
+                    error_message = f"{error_message}: {error_detail}"
+                except Exception:
+                    error_message = f"{error_message}: {response.text}"
+
+                self.log.error(error_message)
+
+                if self.debug:
+                    raise Exception(error_message)
+
+                return None
+
+        except Exception as e:
+            self.log.exception(f"Error sending AI event: {e}")
+
+            if self.debug:
+                raise e
+
+            return None
+
     def _enqueue(self, msg, disable_geoip):
         # type: (...) -> Optional[str]
         """Push a new `msg` onto the queue, return `(success, msg)`"""

test_capture_ai_local.py

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