fix: Actually fix LangChain callback in posthoganalytics

Co-authored-by: Peter Kirkham <peter@posthog.com>

.github/workflows/ci.yml

@@ -18,7 +18,7 @@ jobs:
               with:
                   python-version: 3.8
 
-            - uses: actions/cache@v1
+            - uses: actions/cache@v3
               with:
                   path: ~/.cache/pip
                   key: ${{ runner.os }}-pip-${{ hashFiles('setup.py') }}
@@ -33,10 +33,10 @@ jobs:
             - name: Check formatting with black
               run: |
                   black --check .
-            
+
             - name: Lint with flake8
               run: |
-                  flake8 posthog --ignore E501
+                  flake8 posthog --ignore E501,W503
 
             - name: Check import order with isort
               run: |
@@ -47,14 +47,14 @@ jobs:
         runs-on: ubuntu-latest
 
         steps:
-            - uses: actions/checkout@v1
+            - uses: actions/checkout@v2
               with:
                   fetch-depth: 1
 
-            - name: Set up Python 3.7
-              uses: actions/setup-python@v1
+            - name: Set up Python 3.9
+              uses: actions/setup-python@v2
               with:
-                  python-version: 3.7
+                  python-version: 3.9
 
             - name: Install requirements.txt dependencies with pip
               run: |

.gitignore

@@ -14,4 +14,6 @@ pylint.out
 posthog-analytics
 .idea
 .python-version
-.coverage
+.coverage
+pyrightconfig.json
+.env

CHANGELOG.md

@@ -1,3 +1,48 @@
+## 3.9.2 - 2025-01-22
+
+1. Fix importing of LangChain callback handler under certain circumstances.
+
+## 3.9.0 - 2025-01-22
+
+1. Add `$ai_trace` event emission to LangChain callback handler.
+
+## 3.8.4 - 2025-01-17
+
+1. Add Anthropic support for LLM Observability.
+2. Update LLM Observability to use output_choices.
+
+## 3.8.3 - 2025-01-14
+
+1. Fix setuptools to include the `posthog.ai.openai` and `posthog.ai.langchain` packages for the `posthoganalytics` package.
+
+## 3.8.2 - 2025-01-14
+
+1. Fix setuptools to include the `posthog.ai.openai` and `posthog.ai.langchain` packages.
+
+## 3.8.1 - 2025-01-14
+
+1. Add LLM Observability with support for OpenAI and Langchain callbacks.
+
+## 3.7.5 - 2025-01-03
+
+1. Add `distinct_id` to group_identify
+
+## 3.7.4 - 2024-11-25
+
+1. Fix bug where this SDK incorrectly sent feature flag events with null values when calling `get_feature_flag_payload`.
+
+## 3.7.3 - 2024-11-25
+
+1. Use personless mode when sending an exception without a provided `distinct_id`.
+
+## 3.7.2 - 2024-11-19
+
+1. Add `type` property to exception stacks.
+
+## 3.7.1 - 2024-10-24
+
+1. Add `platform` property to each frame of exception stacks.
+
 ## 3.7.0 - 2024-10-03
 
 1. Adds a new `super_properties` parameter on the client that are appended to every /capture call.

CODEOWNERS

@@ -1 +1 @@
-@PostHog/team-feature-success
+@PostHog/team-feature-flags

Makefile

@@ -17,11 +17,13 @@ release_analytics:
 	rm -rf posthoganalytics
 	mkdir posthoganalytics
 	cp -r posthog/* posthoganalytics/
+	find ./posthoganalytics -type f -exec sed -i '' -e 's/from posthog /from posthoganalytics /g' {} \;
 	find ./posthoganalytics -type f -exec sed -i '' -e 's/from posthog\./from posthoganalytics\./g' {} \;
 	rm -rf posthog
 	python setup_analytics.py sdist bdist_wheel
 	twine upload dist/*
 	mkdir posthog
+	find ./posthoganalytics -type f -exec sed -i '' -e 's/from posthoganalytics /from posthog /g' {} \;
 	find ./posthoganalytics -type f -exec sed -i '' -e 's/from posthoganalytics\./from posthog\./g' {} \;
 	cp -r posthoganalytics/* posthog/
 	rm -rf posthoganalytics

example.py

@@ -1,7 +1,7 @@
 # PostHog Python library example
 
 # Import the library
-import time
+# import time
 
 import posthog
 

llm_observability_examples.py

@@ -0,0 +1,196 @@
+import os
+import uuid
+
+import posthog
+from posthog.ai.openai import AsyncOpenAI, OpenAI
+
+# Example credentials - replace these with your own or use environment variables
+posthog.project_api_key = os.getenv("POSTHOG_PROJECT_API_KEY", "your-project-api-key")
+posthog.personal_api_key = os.getenv("POSTHOG_PERSONAL_API_KEY", "your-personal-api-key")
+posthog.host = os.getenv("POSTHOG_HOST", "http://localhost:8000")  # Or https://app.posthog.com
+posthog.debug = True
+# change this to False to see usage events
+# posthog.privacy_mode = True
+
+openai_client = OpenAI(
+    api_key=os.getenv("OPENAI_API_KEY", "your-openai-api-key"),
+    posthog_client=posthog,
+)
+
+async_openai_client = AsyncOpenAI(
+    api_key=os.getenv("OPENAI_API_KEY", "your-openai-api-key"),
+    posthog_client=posthog,
+)
+
+
+def main_sync():
+    trace_id = str(uuid.uuid4())
+    print("Trace ID:", trace_id)
+    distinct_id = "test2_distinct_id"
+    properties = {"test_property": "test_value"}
+    groups = {"company": "test_company"}
+
+    try:
+        basic_openai_call(distinct_id, trace_id, properties, groups)
+        streaming_openai_call(distinct_id, trace_id, properties, groups)
+        embedding_openai_call(distinct_id, trace_id, properties, groups)
+        image_openai_call()
+    except Exception as e:
+        print("Error during OpenAI call:", str(e))
+
+
+async def main_async():
+    trace_id = str(uuid.uuid4())
+    print("Trace ID:", trace_id)
+    distinct_id = "test_distinct_id"
+    properties = {"test_property": "test_value"}
+    groups = {"company": "test_company"}
+
+    try:
+        await basic_async_openai_call(distinct_id, trace_id, properties, groups)
+        await streaming_async_openai_call(distinct_id, trace_id, properties, groups)
+        await embedding_async_openai_call(distinct_id, trace_id, properties, groups)
+        await image_async_openai_call()
+    except Exception as e:
+        print("Error during OpenAI call:", str(e))
+
+
+def basic_openai_call(distinct_id, trace_id, properties, groups):
+    response = openai_client.chat.completions.create(
+        model="gpt-4o-mini",
+        messages=[
+            {"role": "system", "content": "You are a complex problem solver."},
+            {"role": "user", "content": "Explain quantum computing in simple terms."},
+        ],
+        max_tokens=100,
+        temperature=0.7,
+        posthog_distinct_id=distinct_id,
+        posthog_trace_id=trace_id,
+        posthog_properties=properties,
+        posthog_groups=groups,
+    )
+    print(response)
+    if response and response.choices:
+        print("OpenAI response:", response.choices[0].message.content)
+    else:
+        print("No response or unexpected format returned.")
+    return response
+
+
+async def basic_async_openai_call(distinct_id, trace_id, properties, groups):
+    response = await async_openai_client.chat.completions.create(
+        model="gpt-4o-mini",
+        messages=[
+            {"role": "system", "content": "You are a complex problem solver."},
+            {"role": "user", "content": "Explain quantum computing in simple terms."},
+        ],
+        max_tokens=100,
+        temperature=0.7,
+        posthog_distinct_id=distinct_id,
+        posthog_trace_id=trace_id,
+        posthog_properties=properties,
+        posthog_groups=groups,
+    )
+    if response and hasattr(response, "choices"):
+        print("OpenAI response:", response.choices[0].message.content)
+    else:
+        print("No response or unexpected format returned.")
+    return response
+
+
+def streaming_openai_call(distinct_id, trace_id, properties, groups):
+
+    response = openai_client.chat.completions.create(
+        model="gpt-4o-mini",
+        messages=[
+            {"role": "system", "content": "You are a complex problem solver."},
+            {"role": "user", "content": "Explain quantum computing in simple terms."},
+        ],
+        max_tokens=100,
+        temperature=0.7,
+        stream=True,
+        posthog_distinct_id=distinct_id,
+        posthog_trace_id=trace_id,
+        posthog_properties=properties,
+        posthog_groups=groups,
+    )
+
+    for chunk in response:
+        if hasattr(chunk, "choices") and chunk.choices and len(chunk.choices) > 0:
+            print(chunk.choices[0].delta.content or "", end="")
+
+    return response
+
+
+async def streaming_async_openai_call(distinct_id, trace_id, properties, groups):
+    response = await async_openai_client.chat.completions.create(
+        model="gpt-4o-mini",
+        messages=[
+            {"role": "system", "content": "You are a complex problem solver."},
+            {"role": "user", "content": "Explain quantum computing in simple terms."},
+        ],
+        max_tokens=100,
+        temperature=0.7,
+        stream=True,
+        posthog_distinct_id=distinct_id,
+        posthog_trace_id=trace_id,
+        posthog_properties=properties,
+        posthog_groups=groups,
+    )
+
+    async for chunk in response:
+        if hasattr(chunk, "choices") and chunk.choices and len(chunk.choices) > 0:
+            print(chunk.choices[0].delta.content or "", end="")
+
+    return response
+
+
+# none instrumented
+def image_openai_call():
+    response = openai_client.images.generate(model="dall-e-3", prompt="A cute baby hedgehog", n=1, size="1024x1024")
+    print(response)
+    return response
+
+
+# none instrumented
+async def image_async_openai_call():
+    response = await async_openai_client.images.generate(
+        model="dall-e-3", prompt="A cute baby hedgehog", n=1, size="1024x1024"
+    )
+    print(response)
+    return response
+
+
+def embedding_openai_call(posthog_distinct_id, posthog_trace_id, posthog_properties, posthog_groups):
+    response = openai_client.embeddings.create(
+        input="The hedgehog is cute",
+        model="text-embedding-3-small",
+        posthog_distinct_id=posthog_distinct_id,
+        posthog_trace_id=posthog_trace_id,
+        posthog_properties=posthog_properties,
+        posthog_groups=posthog_groups,
+    )
+    print(response)
+    return response
+
+
+async def embedding_async_openai_call(posthog_distinct_id, posthog_trace_id, posthog_properties, posthog_groups):
+    response = await async_openai_client.embeddings.create(
+        input="The hedgehog is cute",
+        model="text-embedding-3-small",
+        posthog_distinct_id=posthog_distinct_id,
+        posthog_trace_id=posthog_trace_id,
+        posthog_properties=posthog_properties,
+        posthog_groups=posthog_groups,
+    )
+    print(response)
+    return response
+
+
+# HOW TO RUN:
+# comment out one of these to run the other
+
+if __name__ == "__main__":
+    main_sync()
+
+# asyncio.run(main_async())

posthog/__init__.py

@@ -2,7 +2,7 @@ import datetime  # noqa: F401
 from typing import Callable, Dict, List, Optional, Tuple  # noqa: F401
 
 from posthog.client import Client
-from posthog.exception_capture import DEFAULT_DISTINCT_ID, Integrations  # noqa: F401
+from posthog.exception_capture import Integrations  # noqa: F401
 from posthog.version import VERSION
 
 __version__ = VERSION
@@ -26,6 +26,8 @@ enable_exception_autocapture = False  # type: bool
 exception_autocapture_integrations = []  # type: List[Integrations]
 # Used to determine in app paths for exception autocapture. Defaults to the current working directory
 project_root = None  # type: Optional[str]
+# Used for our AI observability feature to not capture any prompt or output just usage + metadata
+privacy_mode = False  # type: bool
 
 default_client = None  # type: Optional[Client]
 
@@ -289,7 +291,7 @@ def capture_exception(
     return _proxy(
         "capture_exception",
         exception=exception,
-        distinct_id=distinct_id or DEFAULT_DISTINCT_ID,
+        distinct_id=distinct_id,
         properties=properties,
         context=context,
         timestamp=timestamp,

posthog/ai/anthropic/__init__.py

@@ -0,0 +1,12 @@
+from .anthropic import Anthropic
+from .anthropic_async import AsyncAnthropic
+from .anthropic_providers import AnthropicBedrock, AnthropicVertex, AsyncAnthropicBedrock, AsyncAnthropicVertex
+
+__all__ = [
+    "Anthropic",
+    "AsyncAnthropic",
+    "AnthropicBedrock",
+    "AsyncAnthropicBedrock",
+    "AnthropicVertex",
+    "AsyncAnthropicVertex",
+]

posthog/ai/anthropic/anthropic.py

@@ -0,0 +1,202 @@
+try:
+    import anthropic
+    from anthropic.resources import Messages
+except ImportError:
+    raise ModuleNotFoundError("Please install the Anthropic SDK to use this feature: 'pip install anthropic'")
+
+import time
+import uuid
+from typing import Any, Dict, Optional
+
+from posthog.ai.utils import call_llm_and_track_usage, get_model_params, merge_system_prompt, with_privacy_mode
+from posthog.client import Client as PostHogClient
+
+
+class Anthropic(anthropic.Anthropic):
+    """
+    A wrapper around the Anthropic SDK that automatically sends LLM usage events to PostHog.
+    """
+
+    _ph_client: PostHogClient
+
+    def __init__(self, posthog_client: PostHogClient, **kwargs):
+        """
+        Args:
+            posthog_client: PostHog client for tracking usage
+            **kwargs: Additional arguments passed to the Anthropic client
+        """
+        super().__init__(**kwargs)
+        self._ph_client = posthog_client
+        self.messages = WrappedMessages(self)
+
+
+class WrappedMessages(Messages):
+    _client: Anthropic
+
+    def create(
+        self,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: bool = False,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        """
+        Create a message using Anthropic's API while tracking usage in PostHog.
+
+        Args:
+            posthog_distinct_id: Optional ID to associate with the usage event
+            posthog_trace_id: Optional trace UUID for linking events
+            posthog_properties: Optional dictionary of extra properties to include in the event
+            posthog_privacy_mode: Whether to redact sensitive information in tracking
+            posthog_groups: Optional group analytics properties
+            **kwargs: Arguments passed to Anthropic's messages.create
+        """
+        if posthog_trace_id is None:
+            posthog_trace_id = uuid.uuid4()
+
+        if kwargs.get("stream", False):
+            return self._create_streaming(
+                posthog_distinct_id,
+                posthog_trace_id,
+                posthog_properties,
+                posthog_privacy_mode,
+                posthog_groups,
+                **kwargs,
+            )
+
+        return call_llm_and_track_usage(
+            posthog_distinct_id,
+            self._client._ph_client,
+            "anthropic",
+            posthog_trace_id,
+            posthog_properties,
+            posthog_privacy_mode,
+            posthog_groups,
+            self._client.base_url,
+            super().create,
+            **kwargs,
+        )
+
+    def stream(
+        self,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: bool = False,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        if posthog_trace_id is None:
+            posthog_trace_id = uuid.uuid4()
+
+        return self._create_streaming(
+            posthog_distinct_id,
+            posthog_trace_id,
+            posthog_properties,
+            posthog_privacy_mode,
+            posthog_groups,
+            **kwargs,
+        )
+
+    def _create_streaming(
+        self,
+        posthog_distinct_id: Optional[str],
+        posthog_trace_id: Optional[str],
+        posthog_properties: Optional[Dict[str, Any]],
+        posthog_privacy_mode: bool,
+        posthog_groups: Optional[Dict[str, Any]],
+        **kwargs: Any,
+    ):
+        start_time = time.time()
+        usage_stats: Dict[str, int] = {"input_tokens": 0, "output_tokens": 0}
+        accumulated_content = []
+        response = super().create(**kwargs)
+
+        def generator():
+            nonlocal usage_stats
+            nonlocal accumulated_content
+            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",
+                            ]
+                        }
+
+                    if hasattr(event, "content") and event.content:
+                        accumulated_content.append(event.content)
+
+                    yield event
+
+            finally:
+                end_time = time.time()
+                latency = end_time - start_time
+                output = "".join(accumulated_content)
+
+                self._capture_streaming_event(
+                    posthog_distinct_id,
+                    posthog_trace_id,
+                    posthog_properties,
+                    posthog_privacy_mode,
+                    posthog_groups,
+                    kwargs,
+                    usage_stats,
+                    latency,
+                    output,
+                )
+
+        return generator()
+
+    def _capture_streaming_event(
+        self,
+        posthog_distinct_id: Optional[str],
+        posthog_trace_id: Optional[str],
+        posthog_properties: Optional[Dict[str, Any]],
+        posthog_privacy_mode: bool,
+        posthog_groups: Optional[Dict[str, Any]],
+        kwargs: Dict[str, Any],
+        usage_stats: Dict[str, int],
+        latency: float,
+        output: str,
+    ):
+        if posthog_trace_id is None:
+            posthog_trace_id = uuid.uuid4()
+
+        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_latency": latency,
+            "$ai_trace_id": posthog_trace_id,
+            "$ai_base_url": str(self._client.base_url),
+            **(posthog_properties or {}),
+        }
+
+        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,
+            )

posthog/ai/anthropic/anthropic_async.py

@@ -0,0 +1,202 @@
+try:
+    import anthropic
+    from anthropic.resources import AsyncMessages
+except ImportError:
+    raise ModuleNotFoundError("Please install the Anthropic SDK to use this feature: 'pip install anthropic'")
+
+import time
+import uuid
+from typing import Any, Dict, Optional
+
+from posthog.ai.utils import call_llm_and_track_usage_async, get_model_params, merge_system_prompt, with_privacy_mode
+from posthog.client import Client as PostHogClient
+
+
+class AsyncAnthropic(anthropic.AsyncAnthropic):
+    """
+    An async wrapper around the Anthropic SDK that automatically sends LLM usage events to PostHog.
+    """
+
+    _ph_client: PostHogClient
+
+    def __init__(self, posthog_client: PostHogClient, **kwargs):
+        """
+        Args:
+            posthog_client: PostHog client for tracking usage
+            **kwargs: Additional arguments passed to the Anthropic client
+        """
+        super().__init__(**kwargs)
+        self._ph_client = posthog_client
+        self.messages = AsyncWrappedMessages(self)
+
+
+class AsyncWrappedMessages(AsyncMessages):
+    _client: AsyncAnthropic
+
+    async def create(
+        self,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: bool = False,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        """
+        Create a message using Anthropic's API while tracking usage in PostHog.
+
+        Args:
+            posthog_distinct_id: Optional ID to associate with the usage event
+            posthog_trace_id: Optional trace UUID for linking events
+            posthog_properties: Optional dictionary of extra properties to include in the event
+            posthog_privacy_mode: Whether to redact sensitive information in tracking
+            posthog_groups: Optional group analytics properties
+            **kwargs: Arguments passed to Anthropic's messages.create
+        """
+        if posthog_trace_id is None:
+            posthog_trace_id = uuid.uuid4()
+
+        if kwargs.get("stream", False):
+            return await self._create_streaming(
+                posthog_distinct_id,
+                posthog_trace_id,
+                posthog_properties,
+                posthog_privacy_mode,
+                posthog_groups,
+                **kwargs,
+            )
+
+        return await call_llm_and_track_usage_async(
+            posthog_distinct_id,
+            self._client._ph_client,
+            "anthropic",
+            posthog_trace_id,
+            posthog_properties,
+            posthog_privacy_mode,
+            posthog_groups,
+            self._client.base_url,
+            super().create,
+            **kwargs,
+        )
+
+    async def stream(
+        self,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: bool = False,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        if posthog_trace_id is None:
+            posthog_trace_id = uuid.uuid4()
+
+        return await self._create_streaming(
+            posthog_distinct_id,
+            posthog_trace_id,
+            posthog_properties,
+            posthog_privacy_mode,
+            posthog_groups,
+            **kwargs,
+        )
+
+    async def _create_streaming(
+        self,
+        posthog_distinct_id: Optional[str],
+        posthog_trace_id: Optional[str],
+        posthog_properties: Optional[Dict[str, Any]],
+        posthog_privacy_mode: bool,
+        posthog_groups: Optional[Dict[str, Any]],
+        **kwargs: Any,
+    ):
+        start_time = time.time()
+        usage_stats: Dict[str, int] = {"input_tokens": 0, "output_tokens": 0}
+        accumulated_content = []
+        response = await super().create(**kwargs)
+
+        async def generator():
+            nonlocal usage_stats
+            nonlocal accumulated_content
+            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",
+                            ]
+                        }
+
+                    if hasattr(event, "content") and event.content:
+                        accumulated_content.append(event.content)
+
+                    yield event
+
+            finally:
+                end_time = time.time()
+                latency = end_time - start_time
+                output = "".join(accumulated_content)
+
+                await self._capture_streaming_event(
+                    posthog_distinct_id,
+                    posthog_trace_id,
+                    posthog_properties,
+                    posthog_privacy_mode,
+                    posthog_groups,
+                    kwargs,
+                    usage_stats,
+                    latency,
+                    output,
+                )
+
+        return generator()
+
+    async def _capture_streaming_event(
+        self,
+        posthog_distinct_id: Optional[str],
+        posthog_trace_id: Optional[str],
+        posthog_properties: Optional[Dict[str, Any]],
+        posthog_privacy_mode: bool,
+        posthog_groups: Optional[Dict[str, Any]],
+        kwargs: Dict[str, Any],
+        usage_stats: Dict[str, int],
+        latency: float,
+        output: str,
+    ):
+        if posthog_trace_id is None:
+            posthog_trace_id = uuid.uuid4()
+
+        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_latency": latency,
+            "$ai_trace_id": posthog_trace_id,
+            "$ai_base_url": str(self._client.base_url),
+            **(posthog_properties or {}),
+        }
+
+        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,
+            )

posthog/ai/anthropic/anthropic_providers.py

@@ -0,0 +1,60 @@
+try:
+    import anthropic
+except ImportError:
+    raise ModuleNotFoundError("Please install the Anthropic SDK to use this feature: 'pip install anthropic'")
+
+from posthog.ai.anthropic.anthropic import WrappedMessages
+from posthog.ai.anthropic.anthropic_async import AsyncWrappedMessages
+from posthog.client import Client as PostHogClient
+
+
+class AnthropicBedrock(anthropic.AnthropicBedrock):
+    """
+    A wrapper around the Anthropic Bedrock SDK that automatically sends LLM usage events to PostHog.
+    """
+
+    _ph_client: PostHogClient
+
+    def __init__(self, posthog_client: PostHogClient, **kwargs):
+        super().__init__(**kwargs)
+        self._ph_client = posthog_client
+        self.messages = WrappedMessages(self)
+
+
+class AsyncAnthropicBedrock(anthropic.AsyncAnthropicBedrock):
+    """
+    A wrapper around the Anthropic Bedrock SDK that automatically sends LLM usage events to PostHog.
+    """
+
+    _ph_client: PostHogClient
+
+    def __init__(self, posthog_client: PostHogClient, **kwargs):
+        super().__init__(**kwargs)
+        self._ph_client = posthog_client
+        self.messages = AsyncWrappedMessages(self)
+
+
+class AnthropicVertex(anthropic.AnthropicVertex):
+    """
+    A wrapper around the Anthropic Vertex SDK that automatically sends LLM usage events to PostHog.
+    """
+
+    _ph_client: PostHogClient
+
+    def __init__(self, posthog_client: PostHogClient, **kwargs):
+        super().__init__(**kwargs)
+        self._ph_client = posthog_client
+        self.messages = WrappedMessages(self)
+
+
+class AsyncAnthropicVertex(anthropic.AsyncAnthropicVertex):
+    """
+    A wrapper around the Anthropic Vertex SDK that automatically sends LLM usage events to PostHog.
+    """
+
+    _ph_client: PostHogClient
+
+    def __init__(self, posthog_client: PostHogClient, **kwargs):
+        super().__init__(**kwargs)
+        self._ph_client = posthog_client
+        self.messages = AsyncWrappedMessages(self)

posthog/ai/langchain/__init__.py

@@ -0,0 +1,3 @@
+from .callbacks import CallbackHandler
+
+__all__ = ["CallbackHandler"]

posthog/ai/langchain/callbacks.py

@@ -0,0 +1,597 @@
+try:
+    import langchain  # noqa: F401
+except ImportError:
+    raise ModuleNotFoundError("Please install LangChain to use this feature: 'pip install langchain'")
+
+import logging
+import time
+import uuid
+from typing import (
+    Any,
+    Dict,
+    List,
+    Optional,
+    Tuple,
+    TypedDict,
+    Union,
+    cast,
+)
+from uuid import UUID
+
+from langchain.callbacks.base import BaseCallbackHandler
+from langchain.schema.agent import AgentAction, AgentFinish
+from langchain_core.messages import AIMessage, BaseMessage, FunctionMessage, HumanMessage, SystemMessage, ToolMessage
+from langchain_core.outputs import ChatGeneration, LLMResult
+from pydantic import BaseModel
+
+from posthog import default_client
+from posthog.ai.utils import get_model_params, with_privacy_mode
+from posthog.client import Client
+
+log = logging.getLogger("posthog")
+
+
+class RunMetadata(TypedDict, total=False):
+    messages: Union[List[Dict[str, Any]], List[str]]
+    provider: str
+    model: str
+    model_params: Dict[str, Any]
+    base_url: str
+    start_time: float
+    end_time: float
+
+
+RunStorage = Dict[UUID, RunMetadata]
+
+
+class CallbackHandler(BaseCallbackHandler):
+    """
+    The PostHog LLM observability callback handler for LangChain.
+    """
+
+    _client: Client
+    """PostHog client instance."""
+
+    _distinct_id: Optional[Union[str, int, float, UUID]]
+    """Distinct ID of the user to associate the trace with."""
+
+    _trace_id: Optional[Union[str, int, float, UUID]]
+    """Global trace ID to be sent with every event. Otherwise, the top-level run ID is used."""
+
+    _trace_input: Optional[Any]
+    """The input at the start of the trace. Any JSON object."""
+
+    _trace_name: Optional[str]
+    """Name of the trace, exposed in the UI."""
+
+    _properties: Optional[Dict[str, Any]]
+    """Global properties to be sent with every event."""
+
+    _runs: RunStorage
+    """Mapping of run IDs to run metadata as run metadata is only available on the start of generation."""
+
+    _parent_tree: Dict[UUID, UUID]
+    """
+    A dictionary that maps chain run IDs to their parent chain run IDs (parent pointer tree),
+    so the top level can be found from a bottom-level run ID.
+    """
+
+    def __init__(
+        self,
+        client: Optional[Client] = None,
+        *,
+        distinct_id: Optional[Union[str, int, float, UUID]] = None,
+        trace_id: Optional[Union[str, int, float, UUID]] = None,
+        properties: Optional[Dict[str, Any]] = None,
+        privacy_mode: bool = False,
+        groups: Optional[Dict[str, Any]] = None,
+    ):
+        """
+        Args:
+            client: PostHog client instance.
+            distinct_id: Optional distinct ID of the user to associate the trace with.
+            trace_id: Optional trace ID to use for the event.
+            properties: Optional additional metadata to use for the trace.
+            privacy_mode: Whether to redact the input and output of the trace.
+            groups: Optional additional PostHog groups to use for the trace.
+        """
+        self._client = client or default_client
+        self._distinct_id = distinct_id
+        self._trace_id = trace_id
+        self._trace_name = None
+        self._trace_input = None
+        self._properties = properties or {}
+        self._privacy_mode = privacy_mode
+        self._groups = groups or {}
+        self._runs = {}
+        self._parent_tree = {}
+
+    def on_chain_start(
+        self,
+        serialized: Dict[str, Any],
+        inputs: Dict[str, Any],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        **kwargs,
+    ):
+        self._log_debug_event("on_chain_start", run_id, parent_run_id, inputs=inputs)
+        self._set_parent_of_run(run_id, parent_run_id)
+        if parent_run_id is None and self._trace_name is None:
+            self._trace_name = self._get_langchain_run_name(serialized, **kwargs)
+            self._trace_input = inputs
+
+    def on_chat_model_start(
+        self,
+        serialized: Dict[str, Any],
+        messages: List[List[BaseMessage]],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs,
+    ):
+        self._log_debug_event("on_chat_model_start", run_id, parent_run_id, messages=messages)
+        self._set_parent_of_run(run_id, parent_run_id)
+        input = [_convert_message_to_dict(message) for row in messages for message in row]
+        self._set_run_metadata(serialized, run_id, input, **kwargs)
+
+    def on_llm_start(
+        self,
+        serialized: Dict[str, Any],
+        prompts: List[str],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ):
+        self._log_debug_event("on_llm_start", run_id, parent_run_id, prompts=prompts)
+        self._set_parent_of_run(run_id, parent_run_id)
+        self._set_run_metadata(serialized, run_id, prompts, **kwargs)
+
+    def on_llm_new_token(
+        self,
+        token: str,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Run on new LLM token. Only available when streaming is enabled."""
+        self._log_debug_event("on_llm_new_token", run_id, parent_run_id, token=token)
+
+    def on_tool_start(
+        self,
+        serialized: Optional[Dict[str, Any]],
+        input_str: str,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> Any:
+        self._log_debug_event("on_tool_start", run_id, parent_run_id, input_str=input_str)
+
+    def on_tool_end(
+        self,
+        output: str,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> Any:
+        self._log_debug_event("on_tool_end", run_id, parent_run_id, output=output)
+
+    def on_tool_error(
+        self,
+        error: Union[Exception, KeyboardInterrupt],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> Any:
+        self._log_debug_event("on_tool_error", run_id, parent_run_id, error=error)
+
+    def on_chain_end(
+        self,
+        outputs: Dict[str, Any],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ):
+        self._log_debug_event("on_chain_end", run_id, parent_run_id, outputs=outputs)
+        self._pop_parent_of_run(run_id)
+
+        if parent_run_id is None:
+            self._capture_trace(run_id, outputs=outputs)
+
+    def on_chain_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ):
+        self._log_debug_event("on_chain_error", run_id, parent_run_id, error=error)
+        self._pop_parent_of_run(run_id)
+
+        if parent_run_id is None:
+            self._capture_trace(run_id, outputs=None)
+
+    def on_llm_end(
+        self,
+        response: LLMResult,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ):
+        """
+        The callback works for both streaming and non-streaming runs. For streaming runs, the chain must set `stream_usage=True` in the LLM.
+        """
+        self._log_debug_event("on_llm_end", run_id, parent_run_id, response=response, kwargs=kwargs)
+        trace_id = self._get_trace_id(run_id)
+        self._pop_parent_of_run(run_id)
+        run = self._pop_run_metadata(run_id)
+        if not run:
+            return
+
+        latency = run.get("end_time", 0) - run.get("start_time", 0)
+        input_tokens, output_tokens = _parse_usage(response)
+
+        generation_result = response.generations[-1]
+        if isinstance(generation_result[-1], ChatGeneration):
+            output = [
+                _convert_message_to_dict(cast(ChatGeneration, generation).message) for generation in generation_result
+            ]
+        else:
+            output = [_extract_raw_esponse(generation) for generation in generation_result]
+
+        event_properties = {
+            "$ai_provider": run.get("provider"),
+            "$ai_model": run.get("model"),
+            "$ai_model_parameters": run.get("model_params"),
+            "$ai_input": with_privacy_mode(self._client, self._privacy_mode, run.get("messages")),
+            "$ai_output_choices": with_privacy_mode(self._client, self._privacy_mode, output),
+            "$ai_http_status": 200,
+            "$ai_input_tokens": input_tokens,
+            "$ai_output_tokens": output_tokens,
+            "$ai_latency": latency,
+            "$ai_trace_id": trace_id,
+            "$ai_base_url": run.get("base_url"),
+            **self._properties,
+        }
+        if self._distinct_id is None:
+            event_properties["$process_person_profile"] = False
+        self._client.capture(
+            distinct_id=self._distinct_id or trace_id,
+            event="$ai_generation",
+            properties=event_properties,
+            groups=self._groups,
+        )
+
+    def on_llm_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ):
+        self._log_debug_event("on_llm_error", run_id, parent_run_id, error=error)
+        trace_id = self._get_trace_id(run_id)
+        self._pop_parent_of_run(run_id)
+        run = self._pop_run_metadata(run_id)
+        if not run:
+            return
+
+        latency = run.get("end_time", 0) - run.get("start_time", 0)
+        event_properties = {
+            "$ai_provider": run.get("provider"),
+            "$ai_model": run.get("model"),
+            "$ai_model_parameters": run.get("model_params"),
+            "$ai_input": with_privacy_mode(self._client, self._privacy_mode, run.get("messages")),
+            "$ai_http_status": _get_http_status(error),
+            "$ai_latency": latency,
+            "$ai_trace_id": trace_id,
+            "$ai_base_url": run.get("base_url"),
+            **self._properties,
+        }
+        if self._distinct_id is None:
+            event_properties["$process_person_profile"] = False
+        self._client.capture(
+            distinct_id=self._distinct_id or trace_id,
+            event="$ai_generation",
+            properties=event_properties,
+            groups=self._groups,
+        )
+
+    def on_retriever_start(
+        self,
+        serialized: Optional[Dict[str, Any]],
+        query: str,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> Any:
+        self._log_debug_event("on_retriever_start", run_id, parent_run_id, query=query)
+
+    def on_retriever_error(
+        self,
+        error: Union[Exception, KeyboardInterrupt],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Run when Retriever errors."""
+        self._log_debug_event("on_retriever_error", run_id, parent_run_id, error=error)
+
+    def on_agent_action(
+        self,
+        action: AgentAction,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Run on agent action."""
+        self._log_debug_event("on_agent_action", run_id, parent_run_id, action=action)
+
+    def on_agent_finish(
+        self,
+        finish: AgentFinish,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> Any:
+        self._log_debug_event("on_agent_finish", run_id, parent_run_id, finish=finish)
+
+    def _set_parent_of_run(self, run_id: UUID, parent_run_id: Optional[UUID] = None):
+        """
+        Set the parent run ID for a chain run. If there is no parent, the run is the root.
+        """
+        if parent_run_id is not None:
+            self._parent_tree[run_id] = parent_run_id
+
+    def _pop_parent_of_run(self, run_id: UUID):
+        """
+        Remove the parent run ID for a chain run.
+        """
+        try:
+            self._parent_tree.pop(run_id)
+        except KeyError:
+            pass
+
+    def _find_root_run(self, run_id: UUID) -> UUID:
+        """
+        Finds the root ID of a chain run.
+        """
+        id: UUID = run_id
+        while id in self._parent_tree:
+            id = self._parent_tree[id]
+        return id
+
+    def _set_run_metadata(
+        self,
+        serialized: Dict[str, Any],
+        run_id: UUID,
+        messages: Union[List[Dict[str, Any]], List[str]],
+        metadata: Optional[Dict[str, Any]] = None,
+        invocation_params: Optional[Dict[str, Any]] = None,
+        **kwargs,
+    ):
+        run: RunMetadata = {
+            "messages": messages,
+            "start_time": time.time(),
+        }
+        if isinstance(invocation_params, dict):
+            run["model_params"] = get_model_params(invocation_params)
+        if isinstance(metadata, dict):
+            if model := metadata.get("ls_model_name"):
+                run["model"] = model
+            if provider := metadata.get("ls_provider"):
+                run["provider"] = provider
+        try:
+            base_url = serialized["kwargs"]["openai_api_base"]
+            if base_url is not None:
+                run["base_url"] = base_url
+        except KeyError:
+            pass
+        self._runs[run_id] = run
+
+    def _pop_run_metadata(self, run_id: UUID) -> Optional[RunMetadata]:
+        end_time = time.time()
+        try:
+            run = self._runs.pop(run_id)
+        except KeyError:
+            log.warning(f"No run metadata found for run {run_id}")
+            return None
+        run["end_time"] = end_time
+        return run
+
+    def _get_trace_id(self, run_id: UUID):
+        trace_id = self._trace_id or self._find_root_run(run_id)
+        if not trace_id:
+            trace_id = uuid.uuid4()
+        return trace_id
+
+    def _get_langchain_run_name(self, serialized: Optional[Dict[str, Any]], **kwargs: Any) -> str:
+        """Retrieve the name of a serialized LangChain runnable.
+
+        The prioritization for the determination of the run name is as follows:
+        - The value assigned to the "name" key in `kwargs`.
+        - The value assigned to the "name" key in `serialized`.
+        - The last entry of the value assigned to the "id" key in `serialized`.
+        - "<unknown>".
+
+        Args:
+            serialized (Optional[Dict[str, Any]]): A dictionary containing the runnable's serialized data.
+            **kwargs (Any): Additional keyword arguments, potentially including the 'name' override.
+
+        Returns:
+            str: The determined name of the Langchain runnable.
+        """
+        if "name" in kwargs and kwargs["name"] is not None:
+            return kwargs["name"]
+
+        try:
+            return serialized["name"]
+        except (KeyError, TypeError):
+            pass
+
+        try:
+            return serialized["id"][-1]
+        except (KeyError, TypeError):
+            pass
+
+    def _capture_trace(self, run_id: UUID, *, outputs: Optional[Dict[str, Any]]):
+        trace_id = self._get_trace_id(run_id)
+        event_properties = {
+            "$ai_trace_name": self._trace_name,
+            "$ai_trace_id": trace_id,
+            "$ai_input_state": with_privacy_mode(self._client, self._privacy_mode, self._trace_input),
+            **self._properties,
+        }
+        if outputs is not None:
+            event_properties["$ai_output_state"] = with_privacy_mode(self._client, self._privacy_mode, outputs)
+        if self._distinct_id is None:
+            event_properties["$process_person_profile"] = False
+        self._client.capture(
+            distinct_id=self._distinct_id or trace_id,
+            event="$ai_trace",
+            properties=event_properties,
+            groups=self._groups,
+        )
+
+    def _log_debug_event(
+        self,
+        event_name: str,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs,
+    ):
+        log.debug(
+            f"Event: {event_name}, run_id: {str(run_id)[:5]}, parent_run_id: {str(parent_run_id)[:5]}, kwargs: {kwargs}"
+        )
+
+
+def _extract_raw_esponse(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() != "":
+        return last_response.text.strip()
+    elif hasattr(last_response, "message"):
+        # Additional kwargs contains the response in case of tool usage
+        return last_response.message.additional_kwargs
+    else:
+        # Not tool usage, some LLM responses can be simply empty
+        return ""
+
+
+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}
+    elif isinstance(message, SystemMessage):
+        message_dict = {"role": "system", "content": message.content}
+    elif isinstance(message, ToolMessage):
+        message_dict = {"role": "tool", "content": message.content}
+    elif isinstance(message, FunctionMessage):
+        message_dict = {"role": "function", "content": message.content}
+    else:
+        message_dict = {"role": message.type, "content": str(message.content)}
+
+    if message.additional_kwargs:
+        message_dict.update(message.additional_kwargs)
+
+    return message_dict
+
+
+def _parse_usage_model(
+    usage: Union[BaseModel, Dict],
+) -> Tuple[Union[int, None], Union[int, None]]:
+    if isinstance(usage, BaseModel):
+        usage = usage.__dict__
+
+    conversion_list = [
+        # https://pypi.org/project/langchain-anthropic/ (works also for Bedrock-Anthropic)
+        ("input_tokens", "input"),
+        ("output_tokens", "output"),
+        # https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/get-token-count
+        ("prompt_token_count", "input"),
+        ("candidates_token_count", "output"),
+        # Bedrock: https://docs.aws.amazon.com/bedrock/latest/userguide/monitoring-cw.html#runtime-cloudwatch-metrics
+        ("inputTokenCount", "input"),
+        ("outputTokenCount", "output"),
+        # langchain-ibm https://pypi.org/project/langchain-ibm/
+        ("input_token_count", "input"),
+        ("generated_token_count", "output"),
+    ]
+
+    parsed_usage = {}
+    for model_key, type_key in conversion_list:
+        if model_key in usage:
+            captured_count = usage[model_key]
+            final_count = (
+                sum(captured_count) if isinstance(captured_count, list) else captured_count
+            )  # For Bedrock, the token count is a list when streamed
+
+            parsed_usage[type_key] = final_count
+
+    return parsed_usage.get("input"), parsed_usage.get("output")
+
+
+def _parse_usage(response: LLMResult):
+    # langchain-anthropic uses the usage field
+    llm_usage_keys = ["token_usage", "usage"]
+    llm_usage: Tuple[Union[int, None], Union[int, None]] = (None, None)
+    if response.llm_output is not None:
+        for key in llm_usage_keys:
+            if response.llm_output.get(key):
+                llm_usage = _parse_usage_model(response.llm_output[key])
+                break
+
+    if hasattr(response, "generations"):
+        for generation in response.generations:
+            for generation_chunk in generation:
+                if generation_chunk.generation_info and ("usage_metadata" in generation_chunk.generation_info):
+                    llm_usage = _parse_usage_model(generation_chunk.generation_info["usage_metadata"])
+                    break
+
+                message_chunk = getattr(generation_chunk, "message", {})
+                response_metadata = getattr(message_chunk, "response_metadata", {})
+
+                bedrock_anthropic_usage = (
+                    response_metadata.get("usage", None)  # for Bedrock-Anthropic
+                    if isinstance(response_metadata, dict)
+                    else None
+                )
+                bedrock_titan_usage = (
+                    response_metadata.get("amazon-bedrock-invocationMetrics", None)  # for Bedrock-Titan
+                    if isinstance(response_metadata, dict)
+                    else None
+                )
+                ollama_usage = getattr(message_chunk, "usage_metadata", None)  # for Ollama
+
+                chunk_usage = bedrock_anthropic_usage or bedrock_titan_usage or ollama_usage
+                if chunk_usage:
+                    llm_usage = _parse_usage_model(chunk_usage)
+                    break
+
+    return llm_usage
+
+
+def _get_http_status(error: BaseException) -> int:
+    # OpenAI: https://github.com/openai/openai-python/blob/main/src/openai/_exceptions.py
+    # Anthropic: https://github.com/anthropics/anthropic-sdk-python/blob/main/src/anthropic/_exceptions.py
+    # Google: https://github.com/googleapis/python-api-core/blob/main/google/api_core/exceptions.py
+    status_code = getattr(error, "status_code", getattr(error, "code", 0))
+    return status_code

posthog/ai/openai/__init__.py

@@ -0,0 +1,4 @@
+from .openai import OpenAI
+from .openai_async import AsyncOpenAI
+
+__all__ = ["OpenAI", "AsyncOpenAI"]

posthog/ai/openai/openai.py

@@ -0,0 +1,251 @@
+import time
+import uuid
+from typing import Any, Dict, Optional
+
+try:
+    import openai
+    import openai.resources
+except ImportError:
+    raise ModuleNotFoundError("Please install the OpenAI SDK to use this feature: 'pip install openai'")
+
+from posthog.ai.utils import call_llm_and_track_usage, get_model_params, with_privacy_mode
+from posthog.client import Client as PostHogClient
+
+
+class OpenAI(openai.OpenAI):
+    """
+    A wrapper around the OpenAI SDK that automatically sends LLM usage events to PostHog.
+    """
+
+    _ph_client: PostHogClient
+
+    def __init__(self, posthog_client: PostHogClient, **kwargs):
+        """
+        Args:
+            api_key: OpenAI API key.
+            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
+        self.chat = WrappedChat(self)
+        self.embeddings = WrappedEmbeddings(self)
+
+
+class WrappedChat(openai.resources.chat.Chat):
+    _client: OpenAI
+
+    @property
+    def completions(self):
+        return WrappedCompletions(self._client)
+
+
+class WrappedCompletions(openai.resources.chat.completions.Completions):
+    _client: OpenAI
+
+    def create(
+        self,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: bool = False,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        if posthog_trace_id is None:
+            posthog_trace_id = uuid.uuid4()
+
+        if kwargs.get("stream", False):
+            return self._create_streaming(
+                posthog_distinct_id,
+                posthog_trace_id,
+                posthog_properties,
+                posthog_privacy_mode,
+                posthog_groups,
+                **kwargs,
+            )
+
+        return call_llm_and_track_usage(
+            posthog_distinct_id,
+            self._client._ph_client,
+            "openai",
+            posthog_trace_id,
+            posthog_properties,
+            posthog_privacy_mode,
+            posthog_groups,
+            self._client.base_url,
+            super().create,
+            **kwargs,
+        )
+
+    def _create_streaming(
+        self,
+        posthog_distinct_id: Optional[str],
+        posthog_trace_id: Optional[str],
+        posthog_properties: Optional[Dict[str, Any]],
+        posthog_privacy_mode: bool,
+        posthog_groups: Optional[Dict[str, Any]],
+        **kwargs: Any,
+    ):
+        start_time = time.time()
+        usage_stats: Dict[str, int] = {}
+        accumulated_content = []
+        if "stream_options" not in kwargs:
+            kwargs["stream_options"] = {}
+        kwargs["stream_options"]["include_usage"] = True
+        response = super().create(**kwargs)
+
+        def generator():
+            nonlocal usage_stats
+            nonlocal accumulated_content
+            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",
+                            ]
+                        }
+
+                    if hasattr(chunk, "choices") and chunk.choices and len(chunk.choices) > 0:
+                        content = chunk.choices[0].delta.content
+                        if content:
+                            accumulated_content.append(content)
+
+                    yield chunk
+
+            finally:
+                end_time = time.time()
+                latency = end_time - start_time
+                output = "".join(accumulated_content)
+                self._capture_streaming_event(
+                    posthog_distinct_id,
+                    posthog_trace_id,
+                    posthog_properties,
+                    posthog_privacy_mode,
+                    posthog_groups,
+                    kwargs,
+                    usage_stats,
+                    latency,
+                    output,
+                )
+
+        return generator()
+
+    def _capture_streaming_event(
+        self,
+        posthog_distinct_id: Optional[str],
+        posthog_trace_id: Optional[str],
+        posthog_properties: Optional[Dict[str, Any]],
+        posthog_privacy_mode: bool,
+        posthog_groups: Optional[Dict[str, Any]],
+        kwargs: Dict[str, Any],
+        usage_stats: Dict[str, int],
+        latency: float,
+        output: str,
+    ):
+        if posthog_trace_id is None:
+            posthog_trace_id = uuid.uuid4()
+
+        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_latency": latency,
+            "$ai_trace_id": posthog_trace_id,
+            "$ai_base_url": str(self._client.base_url),
+            **posthog_properties,
+        }
+
+        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,
+            )
+
+
+class WrappedEmbeddings(openai.resources.embeddings.Embeddings):
+    _client: OpenAI
+
+    def create(
+        self,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: bool = False,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        """
+        Create an embedding using OpenAI's 'embeddings.create' method, but also track usage in PostHog.
+
+        Args:
+            posthog_distinct_id: Optional ID to associate with the usage event.
+            posthog_trace_id: Optional trace UUID for linking events.
+            posthog_properties: Optional dictionary of extra properties to include in the event.
+            **kwargs: Any additional parameters for the OpenAI Embeddings API.
+
+        Returns:
+            The response from OpenAI's embeddings.create call.
+        """
+        if posthog_trace_id is None:
+            posthog_trace_id = uuid.uuid4()
+
+        start_time = time.time()
+        response = super().create(**kwargs)
+        end_time = time.time()
+
+        # Extract usage statistics if available
+        usage_stats = {}
+        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),
+            }
+
+        latency = end_time - start_time
+
+        # Build the event properties
+        event_properties = {
+            "$ai_provider": "openai",
+            "$ai_model": kwargs.get("model"),
+            "$ai_input": with_privacy_mode(self._client._ph_client, posthog_privacy_mode, kwargs.get("input")),
+            "$ai_http_status": 200,
+            "$ai_input_tokens": usage_stats.get("prompt_tokens", 0),
+            "$ai_latency": latency,
+            "$ai_trace_id": posthog_trace_id,
+            "$ai_base_url": str(self._client.base_url),
+            **posthog_properties,
+        }
+
+        if posthog_distinct_id is None:
+            event_properties["$process_person_profile"] = False
+
+        # Send capture event for embeddings
+        if hasattr(self._client._ph_client, "capture"):
+            self._client._ph_client.capture(
+                distinct_id=posthog_distinct_id or posthog_trace_id,
+                event="$ai_embedding",
+                properties=event_properties,
+                groups=posthog_groups,
+            )
+
+        return response

posthog/ai/openai/openai_async.py

@@ -0,0 +1,250 @@
+import time
+import uuid
+from typing import Any, Dict, Optional
+
+try:
+    import openai
+    import openai.resources
+except ImportError:
+    raise ModuleNotFoundError("Please install the OpenAI SDK to use this feature: 'pip install openai'")
+
+from posthog.ai.utils import call_llm_and_track_usage_async, get_model_params, with_privacy_mode
+from posthog.client import Client as PostHogClient
+
+
+class AsyncOpenAI(openai.AsyncOpenAI):
+    """
+    An async wrapper around the OpenAI SDK that automatically sends LLM usage events to PostHog.
+    """
+
+    _ph_client: PostHogClient
+
+    def __init__(self, posthog_client: PostHogClient, **kwargs):
+        """
+        Args:
+            api_key: OpenAI API key.
+            posthog_client: If provided, events will be captured via this client instance.
+            **openai_config: Additional keyword args (e.g. organization="xxx").
+        """
+        super().__init__(**kwargs)
+        self._ph_client = posthog_client
+        self.chat = WrappedChat(self)
+        self.embeddings = WrappedEmbeddings(self)
+
+
+class WrappedChat(openai.resources.chat.AsyncChat):
+    _client: AsyncOpenAI
+
+    @property
+    def completions(self):
+        return WrappedCompletions(self._client)
+
+
+class WrappedCompletions(openai.resources.chat.completions.AsyncCompletions):
+    _client: AsyncOpenAI
+
+    async def create(
+        self,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: bool = False,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        if posthog_trace_id is None:
+            posthog_trace_id = uuid.uuid4()
+
+        # If streaming, handle streaming specifically
+        if kwargs.get("stream", False):
+            return await self._create_streaming(
+                posthog_distinct_id,
+                posthog_trace_id,
+                posthog_properties,
+                posthog_privacy_mode,
+                posthog_groups,
+                **kwargs,
+            )
+
+        response = await call_llm_and_track_usage_async(
+            posthog_distinct_id,
+            self._client._ph_client,
+            "openai",
+            posthog_trace_id,
+            posthog_properties,
+            self._client.base_url,
+            super().create,
+            **kwargs,
+        )
+        return response
+
+    async def _create_streaming(
+        self,
+        posthog_distinct_id: Optional[str],
+        posthog_trace_id: Optional[str],
+        posthog_properties: Optional[Dict[str, Any]],
+        posthog_privacy_mode: bool = False,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        start_time = time.time()
+        usage_stats: Dict[str, int] = {}
+        accumulated_content = []
+        if "stream_options" not in kwargs:
+            kwargs["stream_options"] = {}
+        kwargs["stream_options"]["include_usage"] = True
+        response = await super().create(**kwargs)
+
+        async def async_generator():
+            nonlocal usage_stats, accumulated_content
+            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",
+                            ]
+                        }
+                    if hasattr(chunk, "choices") and chunk.choices and len(chunk.choices) > 0:
+                        content = chunk.choices[0].delta.content
+                        if content:
+                            accumulated_content.append(content)
+
+                    yield chunk
+
+            finally:
+                end_time = time.time()
+                latency = end_time - start_time
+                output = "".join(accumulated_content)
+                await self._capture_streaming_event(
+                    posthog_distinct_id,
+                    posthog_trace_id,
+                    posthog_properties,
+                    posthog_privacy_mode,
+                    posthog_groups,
+                    kwargs,
+                    usage_stats,
+                    latency,
+                    output,
+                )
+
+        return async_generator()
+
+    async def _capture_streaming_event(
+        self,
+        posthog_distinct_id: Optional[str],
+        posthog_trace_id: Optional[str],
+        posthog_properties: Optional[Dict[str, Any]],
+        posthog_privacy_mode: bool,
+        posthog_groups: Optional[Dict[str, Any]],
+        kwargs: Dict[str, Any],
+        usage_stats: Dict[str, int],
+        latency: float,
+        output: str,
+    ):
+        if posthog_trace_id is None:
+            posthog_trace_id = uuid.uuid4()
+
+        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_latency": latency,
+            "$ai_trace_id": posthog_trace_id,
+            "$ai_base_url": str(self._client.base_url),
+            **posthog_properties,
+        }
+
+        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,
+            )
+
+
+class WrappedEmbeddings(openai.resources.embeddings.AsyncEmbeddings):
+    _client: AsyncOpenAI
+
+    async def create(
+        self,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: bool = False,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        """
+        Create an embedding using OpenAI's 'embeddings.create' method, but also track usage in PostHog.
+
+        Args:
+            posthog_distinct_id: Optional ID to associate with the usage event.
+            posthog_trace_id: Optional trace UUID for linking events.
+            posthog_properties: Optional dictionary of extra properties to include in the event.
+            posthog_privacy_mode: Whether to store input and output in PostHog.
+            posthog_groups: Optional dictionary of groups to include in the event.
+            **kwargs: Any additional parameters for the OpenAI Embeddings API.
+
+        Returns:
+            The response from OpenAI's embeddings.create call.
+        """
+        if posthog_trace_id is None:
+            posthog_trace_id = uuid.uuid4()
+
+        start_time = time.time()
+        response = await super().create(**kwargs)
+        end_time = time.time()
+
+        # Extract usage statistics if available
+        usage_stats = {}
+        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),
+            }
+
+        latency = end_time - start_time
+
+        # Build the event properties
+        event_properties = {
+            "$ai_provider": "openai",
+            "$ai_model": kwargs.get("model"),
+            "$ai_input": with_privacy_mode(self._client._ph_client, posthog_privacy_mode, kwargs.get("input")),
+            "$ai_http_status": 200,
+            "$ai_input_tokens": usage_stats.get("prompt_tokens", 0),
+            "$ai_latency": latency,
+            "$ai_trace_id": posthog_trace_id,
+            "$ai_base_url": str(self._client.base_url),
+            **posthog_properties,
+        }
+
+        if posthog_distinct_id is None:
+            event_properties["$process_person_profile"] = False
+
+        # Send capture event for embeddings
+        if hasattr(self._client._ph_client, "capture"):
+            self._client._ph_client.capture(
+                distinct_id=posthog_distinct_id or posthog_trace_id,
+                event="$ai_embedding",
+                properties=event_properties,
+                groups=posthog_groups,
+            )
+
+        return response

posthog/ai/utils.py

@@ -0,0 +1,245 @@
+import time
+import uuid
+from typing import Any, Callable, Dict, Optional
+
+from httpx import URL
+
+from posthog.client import Client as PostHogClient
+
+
+def get_model_params(kwargs: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Extracts model parameters from the kwargs dictionary.
+    """
+    model_params = {}
+    for param in [
+        "temperature",
+        "max_tokens",  # Deprecated field
+        "max_completion_tokens",
+        "top_p",
+        "frequency_penalty",
+        "presence_penalty",
+        "n",
+        "stop",
+        "stream",  # OpenAI-specific field
+        "streaming",  # Anthropic-specific field
+    ]:
+        if param in kwargs and kwargs[param] is not None:
+            model_params[param] = kwargs[param]
+    return model_params
+
+
+def get_usage(response, provider: str) -> Dict[str, Any]:
+    if provider == "anthropic":
+        return {
+            "input_tokens": response.usage.input_tokens,
+            "output_tokens": response.usage.output_tokens,
+        }
+    elif provider == "openai":
+        return {
+            "input_tokens": response.usage.prompt_tokens,
+            "output_tokens": response.usage.completion_tokens,
+        }
+    return {
+        "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)
+    elif provider == "openai":
+        return format_response_openai(response)
+    return output
+
+
+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 = []
+    for choice in response.choices:
+        if choice.message.content:
+            output.append(
+                {
+                    "content": choice.message.content,
+                    "role": choice.message.role,
+                }
+            )
+    return output
+
+
+def merge_system_prompt(kwargs: Dict[str, Any], provider: str):
+    if provider != "anthropic":
+        return kwargs.get("messages")
+    messages = kwargs.get("messages") or []
+    if kwargs.get("system") is None:
+        return messages
+    return [{"role": "system", "content": kwargs.get("system")}] + messages
+
+
+def call_llm_and_track_usage(
+    posthog_distinct_id: Optional[str],
+    ph_client: PostHogClient,
+    provider: str,
+    posthog_trace_id: Optional[str],
+    posthog_properties: Optional[Dict[str, Any]],
+    posthog_privacy_mode: bool,
+    posthog_groups: Optional[Dict[str, Any]],
+    base_url: URL,
+    call_method: Callable[..., Any],
+    **kwargs: Any,
+) -> Any:
+    """
+    Common usage-tracking logic for both sync and async calls.
+    call_method: the llm call method (e.g. openai.chat.completions.create)
+    """
+    start_time = time.time()
+    response = None
+    error = None
+    http_status = 200
+    usage: Dict[str, Any] = {}
+
+    try:
+        response = call_method(**kwargs)
+    except Exception as exc:
+        error = exc
+        http_status = getattr(exc, "status_code", 0)  # default to 0 becuase its likely an SDK error
+    finally:
+        end_time = time.time()
+        latency = end_time - start_time
+
+        if posthog_trace_id is None:
+            posthog_trace_id = uuid.uuid4()
+
+        if response and hasattr(response, "usage"):
+            usage = get_usage(response, provider)
+
+        messages = merge_system_prompt(kwargs, 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_output_choices": with_privacy_mode(
+                ph_client, posthog_privacy_mode, format_response(response, provider)
+            ),
+            "$ai_http_status": http_status,
+            "$ai_input_tokens": usage.get("input_tokens", 0),
+            "$ai_output_tokens": usage.get("output_tokens", 0),
+            "$ai_latency": latency,
+            "$ai_trace_id": posthog_trace_id,
+            "$ai_base_url": str(base_url),
+            **(posthog_properties or {}),
+        }
+
+        if posthog_distinct_id is None:
+            event_properties["$process_person_profile"] = False
+
+        # send the event to posthog
+        if hasattr(ph_client, "capture") and callable(ph_client.capture):
+            ph_client.capture(
+                distinct_id=posthog_distinct_id or posthog_trace_id,
+                event="$ai_generation",
+                properties=event_properties,
+                groups=posthog_groups,
+            )
+
+    if error:
+        raise error
+
+    return response
+
+
+async def call_llm_and_track_usage_async(
+    posthog_distinct_id: Optional[str],
+    ph_client: PostHogClient,
+    provider: str,
+    posthog_trace_id: Optional[str],
+    posthog_properties: Optional[Dict[str, Any]],
+    posthog_privacy_mode: bool,
+    posthog_groups: Optional[Dict[str, Any]],
+    base_url: URL,
+    call_async_method: Callable[..., Any],
+    **kwargs: Any,
+) -> Any:
+    start_time = time.time()
+    response = None
+    error = None
+    http_status = 200
+    usage: Dict[str, Any] = {}
+
+    try:
+        response = await call_async_method(**kwargs)
+    except Exception as exc:
+        error = exc
+        http_status = getattr(exc, "status_code", 0)  # default to 0 because its likely an SDK error
+    finally:
+        end_time = time.time()
+        latency = end_time - start_time
+
+        if posthog_trace_id is None:
+            posthog_trace_id = uuid.uuid4()
+
+        if response and hasattr(response, "usage"):
+            usage = get_usage(response, provider)
+
+        messages = merge_system_prompt(kwargs, 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_output_choices": with_privacy_mode(
+                ph_client, posthog_privacy_mode, format_response(response, provider)
+            ),
+            "$ai_http_status": http_status,
+            "$ai_input_tokens": usage.get("input_tokens", 0),
+            "$ai_output_tokens": usage.get("output_tokens", 0),
+            "$ai_latency": latency,
+            "$ai_trace_id": posthog_trace_id,
+            "$ai_base_url": str(base_url),
+            **(posthog_properties or {}),
+        }
+
+        if posthog_distinct_id is None:
+            event_properties["$process_person_profile"] = False
+
+        # send the event to posthog
+        if hasattr(ph_client, "capture") and callable(ph_client.capture):
+            ph_client.capture(
+                distinct_id=posthog_distinct_id or posthog_trace_id,
+                event="$ai_generation",
+                properties=event_properties,
+                groups=posthog_groups,
+            )
+
+    if error:
+        raise error
+
+    return response
+
+
+def with_privacy_mode(ph_client: PostHogClient, privacy_mode: bool, value: Any):
+    if ph_client.privacy_mode or privacy_mode:
+        return None
+    return value

posthog/client.py

@@ -4,13 +4,13 @@ import numbers
 import os
 import sys
 from datetime import datetime, timedelta
-from uuid import UUID
+from uuid import UUID, uuid4
 
 from dateutil.tz import tzutc
 from six import string_types
 
 from posthog.consumer import Consumer
-from posthog.exception_capture import DEFAULT_DISTINCT_ID, ExceptionCapture
+from posthog.exception_capture import ExceptionCapture
 from posthog.exception_utils import exc_info_from_error, exceptions_from_error_tuple, handle_in_app
 from posthog.feature_flags import InconclusiveMatchError, match_feature_flag_properties
 from posthog.poller import Poller
@@ -59,6 +59,7 @@ class Client(object):
         enable_exception_autocapture=False,
         exception_autocapture_integrations=None,
         project_root=None,
+        privacy_mode=False,
     ):
         self.queue = queue.Queue(max_queue_size)
 
@@ -91,6 +92,7 @@ class Client(object):
         self.enable_exception_autocapture = enable_exception_autocapture
         self.exception_autocapture_integrations = exception_autocapture_integrations
         self.exception_capture = None
+        self.privacy_mode = privacy_mode
 
         if project_root is None:
             try:
@@ -173,6 +175,15 @@ class Client(object):
         resp_data = self.get_decide(distinct_id, groups, person_properties, group_properties, disable_geoip)
         return resp_data["featureFlagPayloads"]
 
+    def get_feature_flags_and_payloads(
+        self, distinct_id, groups=None, person_properties=None, group_properties=None, disable_geoip=None
+    ):
+        resp_data = self.get_decide(distinct_id, groups, person_properties, group_properties, disable_geoip)
+        return {
+            "featureFlags": resp_data["featureFlags"],
+            "featureFlagPayloads": resp_data["featureFlagPayloads"],
+        }
+
     def get_decide(self, distinct_id, groups=None, person_properties=None, group_properties=None, disable_geoip=None):
         require("distinct_id", distinct_id, ID_TYPES)
 
@@ -295,6 +306,7 @@ class Client(object):
         timestamp=None,
         uuid=None,
         disable_geoip=None,
+        distinct_id=None,
     ):
         properties = properties or {}
         context = context or {}
@@ -302,6 +314,11 @@ class Client(object):
         require("group_key", group_key, ID_TYPES)
         require("properties", properties, dict)
 
+        if distinct_id:
+            require("distinct_id", distinct_id, ID_TYPES)
+        else:
+            distinct_id = "${}_{}".format(group_type, group_key)
+
         msg = {
             "event": "$groupidentify",
             "properties": {
@@ -309,7 +326,7 @@ class Client(object):
                 "$group_key": group_key,
                 "$group_set": properties,
             },
-            "distinct_id": "${}_{}".format(group_type, group_key),
+            "distinct_id": distinct_id,
             "timestamp": timestamp,
             "context": context,
             "uuid": uuid,
@@ -362,7 +379,7 @@ class Client(object):
     def capture_exception(
         self,
         exception=None,
-        distinct_id=DEFAULT_DISTINCT_ID,
+        distinct_id=None,
         properties=None,
         context=None,
         timestamp=None,
@@ -373,6 +390,13 @@ class Client(object):
         # this is important to ensure we don't unexpectedly re-raise exceptions in the user's code.
         try:
             properties = properties or {}
+
+            # if there's no distinct_id, we'll generate one and set personless mode
+            # via $process_person_profile = false
+            if distinct_id is None:
+                properties["$process_person_profile"] = False
+                distinct_id = uuid4()
+
             require("distinct_id", distinct_id, ID_TYPES)
             require("properties", properties, dict)
 
@@ -385,7 +409,7 @@ class Client(object):
                 self.log.warning("No exception information available")
                 return
 
-            # Format stack trace like sentry
+            # Format stack trace for cymbal
             all_exceptions_with_trace = exceptions_from_error_tuple(exc_info)
 
             # Add in-app property to frames in the exceptions
@@ -739,23 +763,52 @@ class Client(object):
                 groups=groups,
                 person_properties=person_properties,
                 group_properties=group_properties,
-                send_feature_flag_events=send_feature_flag_events,
-                only_evaluate_locally=True,
+                send_feature_flag_events=False,
+                # Disable automatic sending of feature flag events because we're manually handling event dispatch.
+                # This prevents sending events with empty data when `get_feature_flag` cannot be evaluated locally.
+                only_evaluate_locally=True,  # Enable local evaluation of feature flags to avoid making multiple requests to `/decide`.
                 disable_geoip=disable_geoip,
             )
 
         response = None
+        payload = None
 
         if match_value is not None:
-            response = self._compute_payload_locally(key, match_value)
+            payload = self._compute_payload_locally(key, match_value)
 
-        if response is None and not only_evaluate_locally:
-            decide_payloads = self.get_feature_payloads(
-                distinct_id, groups, person_properties, group_properties, disable_geoip
+        flag_was_locally_evaluated = payload is not None
+        if not flag_was_locally_evaluated and not only_evaluate_locally:
+            try:
+                responses_and_payloads = self.get_feature_flags_and_payloads(
+                    distinct_id, groups, person_properties, group_properties, disable_geoip
+                )
+                response = responses_and_payloads["featureFlags"].get(key, None)
+                payload = responses_and_payloads["featureFlagPayloads"].get(str(key).lower(), None)
+            except Exception as e:
+                self.log.exception(f"[FEATURE FLAGS] Unable to get feature flags and payloads: {e}")
+
+        feature_flag_reported_key = f"{key}_{str(response)}"
+
+        if (
+            feature_flag_reported_key not in self.distinct_ids_feature_flags_reported[distinct_id]
+            and send_feature_flag_events  # noqa: W503
+        ):
+            self.capture(
+                distinct_id,
+                "$feature_flag_called",
+                {
+                    "$feature_flag": key,
+                    "$feature_flag_response": response,
+                    "$feature_flag_payload": payload,
+                    "locally_evaluated": flag_was_locally_evaluated,
+                    f"$feature/{key}": response,
+                },
+                groups=groups,
+                disable_geoip=disable_geoip,
             )
-            response = decide_payloads.get(str(key).lower(), None)
+            self.distinct_ids_feature_flags_reported[distinct_id].add(feature_flag_reported_key)
 
-        return response
+        return payload
 
     def _compute_payload_locally(self, key, match_value):
         payload = None

posthog/exception_capture.py

@@ -12,9 +12,6 @@ class Integrations(str, Enum):
     Django = "django"
 
 
-DEFAULT_DISTINCT_ID = "python-exceptions"
-
-
 class ExceptionCapture:
     # TODO: Add client side rate limiting to prevent spamming the server with exceptions
 
@@ -61,14 +58,7 @@ class ExceptionCapture:
 
     def capture_exception(self, exception, metadata=None):
         try:
-            # if hasattr(sys, "ps1"):
-            #     # Disable the excepthook for interactive Python shells
-            #     return
-
-            distinct_id = metadata.get("distinct_id") if metadata else DEFAULT_DISTINCT_ID
-            # Make sure we have a distinct_id if its empty in metadata
-            distinct_id = distinct_id or DEFAULT_DISTINCT_ID
-
+            distinct_id = metadata.get("distinct_id") if metadata else None
             self.client.capture_exception(exception, distinct_id)
         except Exception as e:
             self.log.exception(f"Failed to capture exception: {e}")

posthog/exception_utils.py

@@ -369,6 +369,7 @@ def serialize_frame(
         tb_lineno = frame.f_lineno
 
     rv = {
+        "platform": "python",
         "filename": filename_for_module(module, abs_path) or None,
         "abs_path": os.path.abspath(abs_path) if abs_path else None,
         "function": function or "<unknown>",
@@ -417,7 +418,7 @@ def current_stacktrace(
 
     frames.reverse()
 
-    return {"frames": frames}
+    return {"frames": frames, "type": "raw"}
 
 
 def get_errno(exc_value):
@@ -503,7 +504,7 @@ def single_exception_from_error_tuple(
     ]
 
     if frames:
-        exception_value["stacktrace"] = {"frames": frames}
+        exception_value["stacktrace"] = {"frames": frames, "type": "raw"}
 
     return exception_value
 

posthog/test/ai/anthropic/test_anthropic.py

@@ -0,0 +1,327 @@
+import os
+import time
+from unittest.mock import patch
+
+import pytest
+from anthropic.types import Message, Usage
+
+from posthog.ai.anthropic import Anthropic, AsyncAnthropic
+
+ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY")
+
+
+@pytest.fixture
+def mock_client():
+    with patch("posthog.client.Client") as mock_client:
+        mock_client.privacy_mode = False
+        yield mock_client
+
+
+@pytest.fixture
+def mock_anthropic_response():
+    return Message(
+        id="msg_123",
+        type="message",
+        role="assistant",
+        content=[{"type": "text", "text": "Test response"}],
+        model="claude-3-opus-20240229",
+        usage=Usage(
+            input_tokens=20,
+            output_tokens=10,
+        ),
+        stop_reason="end_turn",
+        stop_sequence=None,
+    )
+
+
+@pytest.fixture
+def mock_anthropic_stream():
+    class MockStreamEvent:
+        def __init__(self, content, usage=None):
+            self.content = content
+            self.usage = usage
+
+    def stream_generator():
+        yield MockStreamEvent("A")
+        yield MockStreamEvent("B")
+        yield MockStreamEvent(
+            "C",
+            usage=Usage(
+                input_tokens=20,
+                output_tokens=10,
+            ),
+        )
+
+    return stream_generator()
+
+
+def test_basic_completion(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)
+        response = client.messages.create(
+            model="claude-3-opus-20240229",
+            messages=[{"role": "user", "content": "Hello"}],
+            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-opus-20240229"
+        assert props["$ai_input"] == [{"role": "user", "content": "Hello"}]
+        assert props["$ai_output_choices"] == [{"role": "assistant", "content": "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)
+
+
+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):
+        client = Anthropic(api_key="test-key", posthog_client=mock_client)
+        response = client.messages.create(
+            model="claude-3-opus-20240229",
+            messages=[{"role": "user", "content": "Hello"}],
+            posthog_distinct_id="test-id",
+            posthog_groups={"company": "test_company"},
+        )
+
+        assert response == mock_anthropic_response
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+        assert call_args["groups"] == {"company": "test_company"}
+
+
+def test_privacy_mode_local(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)
+        response = client.messages.create(
+            model="claude-3-opus-20240229",
+            messages=[{"role": "user", "content": "Hello"}],
+            posthog_distinct_id="test-id",
+            posthog_privacy_mode=True,
+        )
+
+        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 props["$ai_input"] is None
+        assert props["$ai_output_choices"] is None
+
+
+def test_privacy_mode_global(mock_client, mock_anthropic_response):
+    with patch("anthropic.resources.Messages.create", return_value=mock_anthropic_response):
+        mock_client.privacy_mode = True
+        client = Anthropic(api_key="test-key", posthog_client=mock_client)
+        response = client.messages.create(
+            model="claude-3-opus-20240229",
+            messages=[{"role": "user", "content": "Hello"}],
+            posthog_distinct_id="test-id",
+            posthog_privacy_mode=False,
+        )
+
+        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 props["$ai_input"] is None
+        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'.",
+    )
+
+    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": "system", "content": "You must always answer with 'Bar'."},
+        {"role": "user", "content": "Foo"},
+    ]
+    assert props["$ai_output_choices"][0]["role"] == "assistant"
+    assert props["$ai_output_choices"][0]["content"] == "Bar"
+    assert props["$ai_input_tokens"] == 18
+    assert props["$ai_output_tokens"] == 1
+    assert props["$ai_http_status"] == 200
+    assert props["foo"] == "bar"
+    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"},
+    )
+
+    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": "You must always answer with 'Bar'."}]
+    assert props["$ai_output_choices"][0]["role"] == "assistant"
+    assert props["$ai_input_tokens"] == 16
+    assert props["$ai_output_tokens"] == 1
+    assert props["$ai_http_status"] == 200
+    assert props["foo"] == "bar"
+    assert isinstance(props["$ai_latency"], float)
+
+
+def test_streaming_system_prompt(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",
+            system="Foo",
+            messages=[{"role": "user", "content": "Bar"}],
+            stream=True,
+        )
+
+        # Consume the stream
+        list(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"]
+
+        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"]
+
+    assert props["$ai_input"] == [
+        {"role": "system", "content": "You must always answer with 'Bar'."},
+        {"role": "user", "content": "Foo"},
+    ]

posthog/test/ai/langchain/__init__.py

@@ -0,0 +1,5 @@
+import pytest
+
+pytest.importorskip("langchain")
+pytest.importorskip("langchain_community")
+pytest.importorskip("langgraph")

posthog/test/ai/langchain/test_callbacks.py

@@ -0,0 +1,983 @@
+import logging
+import math
+import os
+import time
+import uuid
+from typing import List, Optional, TypedDict, Union
+from unittest.mock import patch
+
+import pytest
+from langchain_anthropic.chat_models import ChatAnthropic
+from langchain_community.chat_models.fake import FakeMessagesListChatModel
+from langchain_community.llms.fake import FakeListLLM, FakeStreamingListLLM
+from langchain_core.messages import AIMessage, HumanMessage
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.runnables import RunnableLambda
+from langchain_openai.chat_models import ChatOpenAI
+from langgraph.graph.state import END, START, StateGraph
+
+from posthog.ai.langchain import CallbackHandler
+
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY")
+
+
+@pytest.fixture(scope="function")
+def mock_client():
+    with patch("posthog.client.Client") as mock_client:
+        mock_client.privacy_mode = False
+        logging.getLogger("posthog").setLevel(logging.DEBUG)
+        yield mock_client
+
+
+def test_parent_capture(mock_client):
+    callbacks = CallbackHandler(mock_client)
+    parent_run_id = uuid.uuid4()
+    run_id = uuid.uuid4()
+    callbacks._set_parent_of_run(run_id, parent_run_id)
+    assert callbacks._parent_tree == {run_id: parent_run_id}
+    callbacks._pop_parent_of_run(run_id)
+    assert callbacks._parent_tree == {}
+    callbacks._pop_parent_of_run(parent_run_id)  # should not raise
+
+
+def test_find_root_run(mock_client):
+    callbacks = CallbackHandler(mock_client)
+    root_run_id = uuid.uuid4()
+    parent_run_id = uuid.uuid4()
+    run_id = uuid.uuid4()
+    callbacks._set_parent_of_run(run_id, parent_run_id)
+    callbacks._set_parent_of_run(parent_run_id, root_run_id)
+    assert callbacks._find_root_run(run_id) == root_run_id
+    new_run_id = uuid.uuid4()
+    assert callbacks._find_root_run(new_run_id) == new_run_id
+
+
+def test_trace_id_generation(mock_client):
+    callbacks = CallbackHandler(mock_client)
+    run_id = uuid.uuid4()
+    with patch("uuid.uuid4", return_value=run_id):
+        assert callbacks._get_trace_id(run_id) == run_id
+    run_id = uuid.uuid4()
+    callbacks = CallbackHandler(mock_client, trace_id=run_id)
+    assert callbacks._get_trace_id(uuid.uuid4()) == run_id
+
+
+def test_metadata_capture(mock_client):
+    callbacks = CallbackHandler(mock_client)
+    run_id = uuid.uuid4()
+    with patch("time.time", return_value=1234567890):
+        callbacks._set_run_metadata(
+            {"kwargs": {"openai_api_base": "https://us.posthog.com"}},
+            run_id,
+            messages=[{"role": "user", "content": "Who won the world series in 2020?"}],
+            invocation_params={"temperature": 0.5},
+            metadata={"ls_model_name": "hog-mini", "ls_provider": "posthog"},
+        )
+    expected = {
+        "model": "hog-mini",
+        "messages": [{"role": "user", "content": "Who won the world series in 2020?"}],
+        "start_time": 1234567890,
+        "model_params": {"temperature": 0.5},
+        "provider": "posthog",
+        "base_url": "https://us.posthog.com",
+    }
+    assert callbacks._runs[run_id] == expected
+    with patch("time.time", return_value=1234567891):
+        run = callbacks._pop_run_metadata(run_id)
+    assert run == {**expected, "end_time": 1234567891}
+    assert callbacks._runs == {}
+    callbacks._pop_run_metadata(uuid.uuid4())  # should not raise
+
+
+@pytest.mark.parametrize("stream", [True, False])
+def test_basic_chat_chain(mock_client, stream):
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", "You are a helpful assistant."),
+            ("user", "Who won the world series in 2020?"),
+        ]
+    )
+    model = FakeMessagesListChatModel(
+        responses=[
+            AIMessage(
+                content="The Los Angeles Dodgers won the World Series in 2020.",
+                usage_metadata={
+                    "input_tokens": 10,
+                    "output_tokens": 10,
+                    "total_tokens": 20,
+                },
+            )
+        ]
+    )
+    callbacks = [CallbackHandler(mock_client)]
+    chain = prompt | model
+    if stream:
+        result = [m for m in chain.stream({}, config={"callbacks": callbacks})][0]
+    else:
+        result = chain.invoke({}, config={"callbacks": callbacks})
+
+    assert result.content == "The Los Angeles Dodgers won the World Series in 2020."
+    assert mock_client.capture.call_count == 2
+    generation_args = mock_client.capture.call_args_list[0][1]
+    generation_props = generation_args["properties"]
+    trace_args = mock_client.capture.call_args_list[1][1]
+
+    assert generation_args["event"] == "$ai_generation"
+    assert "distinct_id" in generation_args
+    assert "$ai_model" in generation_props
+    assert "$ai_provider" in generation_props
+    assert generation_props["$ai_input"] == [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": "Who won the world series in 2020?"},
+    ]
+    assert generation_props["$ai_output_choices"] == [
+        {
+            "role": "assistant",
+            "content": "The Los Angeles Dodgers won the World Series in 2020.",
+        }
+    ]
+    assert generation_props["$ai_input_tokens"] == 10
+    assert generation_props["$ai_output_tokens"] == 10
+    assert generation_props["$ai_http_status"] == 200
+    assert generation_props["$ai_trace_id"] is not None
+    assert isinstance(generation_props["$ai_latency"], float)
+    assert trace_args["event"] == "$ai_trace"
+
+
+@pytest.mark.parametrize("stream", [True, False])
+async def test_async_basic_chat_chain(mock_client, stream):
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", "You are a helpful assistant."),
+            ("user", "Who won the world series in 2020?"),
+        ]
+    )
+    model = FakeMessagesListChatModel(
+        responses=[
+            AIMessage(
+                content="The Los Angeles Dodgers won the World Series in 2020.",
+                usage_metadata={
+                    "input_tokens": 10,
+                    "output_tokens": 10,
+                    "total_tokens": 20,
+                },
+            )
+        ]
+    )
+    callbacks = [CallbackHandler(mock_client)]
+    chain = prompt | model
+    if stream:
+        result = [m async for m in chain.astream({}, config={"callbacks": callbacks})][0]
+    else:
+        result = await chain.ainvoke({}, config={"callbacks": callbacks})
+    assert result.content == "The Los Angeles Dodgers won the World Series in 2020."
+    assert mock_client.capture.call_count == 2
+
+    generation_args = mock_client.capture.call_args_list[0][1]
+    generation_props = generation_args["properties"]
+    trace_args = mock_client.capture.call_args_list[1][1]
+    trace_props = trace_args["properties"]
+
+    assert generation_args["event"] == "$ai_generation"
+    assert "distinct_id" in generation_args
+    assert "$ai_model" in generation_props
+    assert "$ai_provider" in generation_props
+    assert generation_props["$ai_input"] == [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": "Who won the world series in 2020?"},
+    ]
+    assert generation_props["$ai_output_choices"] == [
+        {
+            "role": "assistant",
+            "content": "The Los Angeles Dodgers won the World Series in 2020.",
+        }
+    ]
+    assert generation_props["$ai_input_tokens"] == 10
+    assert generation_props["$ai_output_tokens"] == 10
+    assert generation_props["$ai_http_status"] == 200
+    assert generation_props["$ai_trace_id"] is not None
+    assert isinstance(generation_props["$ai_latency"], float)
+
+    assert trace_args["event"] == "$ai_trace"
+    assert "distinct_id" in generation_args
+    assert trace_props["$ai_trace_id"] == generation_props["$ai_trace_id"]
+
+
+@pytest.mark.parametrize(
+    "Model,stream",
+    [
+        (FakeListLLM, True),
+        (FakeListLLM, False),
+        (FakeStreamingListLLM, True),
+        (FakeStreamingListLLM, False),
+    ],
+)
+def test_basic_llm_chain(mock_client, Model, stream):
+    model = Model(responses=["The Los Angeles Dodgers won the World Series in 2020."])
+    callbacks: List[CallbackHandler] = [CallbackHandler(mock_client)]
+
+    if stream:
+        result = "".join(
+            [m for m in model.stream("Who won the world series in 2020?", config={"callbacks": callbacks})]
+        )
+    else:
+        result = model.invoke("Who won the world series in 2020?", config={"callbacks": callbacks})
+    assert result == "The Los Angeles Dodgers won the World Series in 2020."
+
+    assert mock_client.capture.call_count == 1
+    args = mock_client.capture.call_args_list[0][1]
+    props = args["properties"]
+
+    assert args["event"] == "$ai_generation"
+    assert "distinct_id" in args
+    assert "$ai_model" in props
+    assert "$ai_provider" in props
+    assert props["$ai_input"] == ["Who won the world series in 2020?"]
+    assert props["$ai_output_choices"] == ["The Los Angeles Dodgers won the World Series in 2020."]
+    assert props["$ai_http_status"] == 200
+    assert props["$ai_trace_id"] is not None
+    assert isinstance(props["$ai_latency"], float)
+
+
+@pytest.mark.parametrize(
+    "Model,stream",
+    [
+        (FakeListLLM, True),
+        (FakeListLLM, False),
+        (FakeStreamingListLLM, True),
+        (FakeStreamingListLLM, False),
+    ],
+)
+async def test_async_basic_llm_chain(mock_client, Model, stream):
+    model = Model(responses=["The Los Angeles Dodgers won the World Series in 2020."])
+    callbacks: List[CallbackHandler] = [CallbackHandler(mock_client)]
+
+    if stream:
+        result = "".join(
+            [m async for m in model.astream("Who won the world series in 2020?", config={"callbacks": callbacks})]
+        )
+    else:
+        result = await model.ainvoke("Who won the world series in 2020?", config={"callbacks": callbacks})
+    assert result == "The Los Angeles Dodgers won the World Series in 2020."
+
+    assert mock_client.capture.call_count == 1
+    args = mock_client.capture.call_args_list[0][1]
+    props = args["properties"]
+
+    assert args["event"] == "$ai_generation"
+    assert "distinct_id" in args
+    assert "$ai_model" in props
+    assert "$ai_provider" in props
+    assert props["$ai_input"] == ["Who won the world series in 2020?"]
+    assert props["$ai_output_choices"] == ["The Los Angeles Dodgers won the World Series in 2020."]
+    assert props["$ai_http_status"] == 200
+    assert props["$ai_trace_id"] is not None
+    assert isinstance(props["$ai_latency"], float)
+
+
+def test_trace_id_for_multiple_chains(mock_client):
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("user", "Foo"),
+        ]
+    )
+    model = FakeMessagesListChatModel(responses=[AIMessage(content="Bar")])
+    callbacks = [CallbackHandler(mock_client)]
+    chain = prompt | model | RunnableLambda(lambda x: [x]) | model
+    result = chain.invoke({}, config={"callbacks": callbacks})
+
+    assert result.content == "Bar"
+    assert mock_client.capture.call_count == 3
+
+    first_call_args = mock_client.capture.call_args_list[0][1]
+    first_call_props = first_call_args["properties"]
+    assert first_call_args["event"] == "$ai_generation"
+    assert "distinct_id" in first_call_args
+    assert "$ai_model" in first_call_props
+    assert "$ai_provider" in first_call_props
+    assert first_call_props["$ai_input"] == [{"role": "user", "content": "Foo"}]
+    assert first_call_props["$ai_output_choices"] == [{"role": "assistant", "content": "Bar"}]
+    assert first_call_props["$ai_http_status"] == 200
+    assert first_call_props["$ai_trace_id"] is not None
+    assert isinstance(first_call_props["$ai_latency"], float)
+
+    second_generation_args = mock_client.capture.call_args_list[1][1]
+    second_generation_props = second_generation_args["properties"]
+    assert second_generation_args["event"] == "$ai_generation"
+    assert "distinct_id" in second_generation_args
+    assert "$ai_model" in second_generation_props
+    assert "$ai_provider" in second_generation_props
+    assert second_generation_props["$ai_input"] == [{"role": "assistant", "content": "Bar"}]
+    assert second_generation_props["$ai_output_choices"] == [{"role": "assistant", "content": "Bar"}]
+    assert second_generation_props["$ai_http_status"] == 200
+    assert second_generation_props["$ai_trace_id"] is not None
+    assert isinstance(second_generation_props["$ai_latency"], float)
+
+    trace_args = mock_client.capture.call_args_list[2][1]
+    trace_props = trace_args["properties"]
+    assert trace_args["event"] == "$ai_trace"
+    assert "distinct_id" in trace_args
+    assert trace_props["$ai_input_state"] == {}
+    assert isinstance(trace_props["$ai_output_state"], AIMessage)
+    assert trace_props["$ai_output_state"].content == "Bar"
+    assert trace_props["$ai_trace_id"] is not None
+    assert trace_props["$ai_trace_name"] == "RunnableSequence"
+
+    # Check that the trace_id is the same as the first call
+    assert first_call_props["$ai_trace_id"] == second_generation_props["$ai_trace_id"]
+    assert first_call_props["$ai_trace_id"] == trace_props["$ai_trace_id"]
+
+
+def test_personless_mode(mock_client):
+    prompt = ChatPromptTemplate.from_messages([("user", "Foo")])
+    chain = prompt | FakeMessagesListChatModel(responses=[AIMessage(content="Bar")])
+    chain.invoke({}, config={"callbacks": [CallbackHandler(mock_client)]})
+    assert mock_client.capture.call_count == 2
+    generation_args = mock_client.capture.call_args_list[0][1]
+    trace_args = mock_client.capture.call_args_list[1][1]
+    assert generation_args["event"] == "$ai_generation"
+    assert generation_args["properties"]["$process_person_profile"] is False
+    assert trace_args["event"] == "$ai_trace"
+    assert trace_args["properties"]["$process_person_profile"] is False
+
+    id = uuid.uuid4()
+    chain.invoke({}, config={"callbacks": [CallbackHandler(mock_client, distinct_id=id)]})
+    assert mock_client.capture.call_count == 4
+    generation_args = mock_client.capture.call_args_list[2][1]
+    trace_args = mock_client.capture.call_args_list[3][1]
+    assert "$process_person_profile" not in generation_args["properties"]
+    assert generation_args["distinct_id"] == id
+    assert "$process_person_profile" not in trace_args["properties"]
+    assert trace_args["distinct_id"] == id
+
+
+def test_personless_mode_exception(mock_client):
+    prompt = ChatPromptTemplate.from_messages([("user", "Foo")])
+    chain = prompt | ChatOpenAI(api_key="test", model="gpt-4o-mini")
+    callbacks = CallbackHandler(mock_client)
+    with pytest.raises(Exception):
+        chain.invoke({}, config={"callbacks": [callbacks]})
+    assert mock_client.capture.call_count == 2
+    generation_args = mock_client.capture.call_args_list[0][1]
+    trace_args = mock_client.capture.call_args_list[1][1]
+    assert generation_args["event"] == "$ai_generation"
+    assert generation_args["properties"]["$process_person_profile"] is False
+    assert trace_args["event"] == "$ai_trace"
+    assert trace_args["properties"]["$process_person_profile"] is False
+
+    id = uuid.uuid4()
+    with pytest.raises(Exception):
+        chain.invoke({}, config={"callbacks": [CallbackHandler(mock_client, distinct_id=id)]})
+    assert mock_client.capture.call_count == 4
+    generation_args = mock_client.capture.call_args_list[2][1]
+    trace_args = mock_client.capture.call_args_list[3][1]
+    assert "$process_person_profile" not in generation_args["properties"]
+    assert generation_args["distinct_id"] == id
+    assert "$process_person_profile" not in trace_args["properties"]
+    assert trace_args["distinct_id"] == id
+
+
+def test_metadata(mock_client):
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("user", "Foo"),
+        ]
+    )
+    model = FakeMessagesListChatModel(responses=[AIMessage(content="Bar")])
+    callbacks = [
+        CallbackHandler(
+            mock_client,
+            trace_id="test-trace-id",
+            distinct_id="test_id",
+            properties={"foo": "bar"},
+        )
+    ]
+    chain = prompt | model
+    result = chain.invoke({"plan": None}, config={"callbacks": callbacks})
+
+    assert result.content == "Bar"
+    assert mock_client.capture.call_count == 2
+
+    generation_call_args = mock_client.capture.call_args_list[0][1]
+    generation_call_props = generation_call_args["properties"]
+    assert generation_call_args["distinct_id"] == "test_id"
+    assert generation_call_args["event"] == "$ai_generation"
+    assert generation_call_props["$ai_trace_id"] == "test-trace-id"
+    assert generation_call_props["foo"] == "bar"
+    assert generation_call_props["$ai_input"] == [{"role": "user", "content": "Foo"}]
+    assert generation_call_props["$ai_output_choices"] == [{"role": "assistant", "content": "Bar"}]
+    assert generation_call_props["$ai_http_status"] == 200
+    assert isinstance(generation_call_props["$ai_latency"], float)
+
+    trace_call_args = mock_client.capture.call_args_list[1][1]
+    trace_call_props = trace_call_args["properties"]
+    assert trace_call_args["distinct_id"] == "test_id"
+    assert trace_call_args["event"] == "$ai_trace"
+    assert trace_call_props["$ai_trace_id"] == "test-trace-id"
+    assert trace_call_props["$ai_trace_name"] == "RunnableSequence"
+    assert trace_call_props["foo"] == "bar"
+    assert trace_call_props["$ai_input_state"] == {"plan": None}
+    assert isinstance(trace_call_props["$ai_output_state"], AIMessage)
+    assert trace_call_props["$ai_output_state"].content == "Bar"
+
+
+class FakeGraphState(TypedDict):
+    messages: List[Union[HumanMessage, AIMessage]]
+    xyz: Optional[str]
+
+
+def test_graph_state(mock_client):
+    config = {"callbacks": [CallbackHandler(mock_client)]}
+
+    graph = StateGraph(FakeGraphState)
+    graph.add_node(
+        "fake_plain",
+        lambda state: {
+            "messages": [
+                *state["messages"],
+                AIMessage(content="Let's explore bar."),
+            ],
+            "xyz": "abc",
+        },
+    )
+    intermediate_chain = ChatPromptTemplate.from_messages(
+        [("user", "Question: What's a bar?")]
+    ) | FakeMessagesListChatModel(
+        responses=[
+            AIMessage(content="It's a type of greeble."),
+        ]
+    )
+    graph.add_node(
+        "fake_llm",
+        lambda state: {
+            "messages": [
+                *state["messages"],
+                intermediate_chain.invoke(state),
+            ],
+            "xyz": state["xyz"],
+        },
+    )
+    graph.add_edge(START, "fake_plain")
+    graph.add_edge("fake_plain", "fake_llm")
+    graph.add_edge("fake_llm", END)
+
+    result = graph.compile().invoke(
+        {"messages": [HumanMessage(content="What's a bar?")], "xyz": None},
+        config=config,
+    )
+
+    assert len(result["messages"]) == 3
+    assert isinstance(result["messages"][0], HumanMessage)
+    assert result["messages"][0].content == "What's a bar?"
+    assert isinstance(result["messages"][1], AIMessage)
+    assert result["messages"][1].content == "Let's explore bar."
+    assert isinstance(result["messages"][2], AIMessage)
+    assert result["messages"][2].content == "It's a type of greeble."
+
+    assert mock_client.capture.call_count == 2
+    generation_args = mock_client.capture.call_args_list[0][1]
+    trace_args = mock_client.capture.call_args_list[1][1]
+    assert generation_args["event"] == "$ai_generation"
+    assert trace_args["event"] == "$ai_trace"
+    assert trace_args["properties"]["$ai_trace_name"] == "LangGraph"
+
+    assert len(trace_args["properties"]["$ai_input_state"]["messages"]) == 1
+    assert isinstance(trace_args["properties"]["$ai_input_state"]["messages"][0], HumanMessage)
+    assert trace_args["properties"]["$ai_input_state"]["messages"][0].content == "What's a bar?"
+    assert trace_args["properties"]["$ai_input_state"]["messages"][0].type == "human"
+    assert trace_args["properties"]["$ai_input_state"]["xyz"] is None
+    assert len(trace_args["properties"]["$ai_output_state"]["messages"]) == 3
+
+    assert isinstance(trace_args["properties"]["$ai_output_state"]["messages"][0], HumanMessage)
+    assert trace_args["properties"]["$ai_output_state"]["messages"][0].content == "What's a bar?"
+    assert isinstance(trace_args["properties"]["$ai_output_state"]["messages"][1], AIMessage)
+    assert trace_args["properties"]["$ai_output_state"]["messages"][1].content == "Let's explore bar."
+    assert isinstance(trace_args["properties"]["$ai_output_state"]["messages"][2], AIMessage)
+    assert trace_args["properties"]["$ai_output_state"]["messages"][2].content == "It's a type of greeble."
+    assert trace_args["properties"]["$ai_output_state"]["xyz"] == "abc"
+
+
+def test_callbacks_logic(mock_client):
+    prompt = ChatPromptTemplate.from_messages([("user", "Foo")])
+    model = FakeMessagesListChatModel(responses=[AIMessage(content="Bar")])
+    callbacks = CallbackHandler(
+        mock_client,
+        trace_id="test-trace-id",
+        distinct_id="test_id",
+        properties={"foo": "bar"},
+    )
+    chain = prompt | model
+
+    chain.invoke({}, config={"callbacks": [callbacks]})
+    assert callbacks._runs == {}
+    assert callbacks._parent_tree == {}
+
+    def assert_intermediary_run(m):
+        assert callbacks._runs == {}
+        assert len(callbacks._parent_tree.items()) == 1
+        return [m]
+
+    (chain | RunnableLambda(assert_intermediary_run) | model).invoke({}, config={"callbacks": [callbacks]})
+    assert callbacks._runs == {}
+    assert callbacks._parent_tree == {}
+
+
+def test_exception_in_chain(mock_client):
+    def runnable(_):
+        raise ValueError("test")
+
+    callbacks = CallbackHandler(mock_client)
+    with pytest.raises(ValueError):
+        RunnableLambda(runnable).invoke({}, config={"callbacks": [callbacks]})
+
+    assert callbacks._runs == {}
+    assert callbacks._parent_tree == {}
+    assert mock_client.capture.call_count == 1
+    trace_call_args = mock_client.capture.call_args_list[0][1]
+    assert trace_call_args["event"] == "$ai_trace"
+    assert trace_call_args["properties"]["$ai_trace_name"] == "runnable"
+
+
+def test_openai_error(mock_client):
+    prompt = ChatPromptTemplate.from_messages([("user", "Foo")])
+    chain = prompt | ChatOpenAI(api_key="test", model="gpt-4o-mini")
+    callbacks = CallbackHandler(mock_client)
+
+    # 401
+    with pytest.raises(Exception):
+        chain.invoke({}, config={"callbacks": [callbacks]})
+
+    assert callbacks._runs == {}
+    assert callbacks._parent_tree == {}
+    assert mock_client.capture.call_count == 2
+    generation_args = mock_client.capture.call_args_list[0][1]
+    props = generation_args["properties"]
+    assert props["$ai_http_status"] == 401
+    assert props["$ai_input"] == [{"role": "user", "content": "Foo"}]
+    assert "$ai_output_choices" not in props
+
+
+@pytest.mark.skipif(not OPENAI_API_KEY, reason="OpenAI API key not set")
+def test_openai_chain(mock_client):
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", 'You must always answer with "Bar".'),
+            ("user", "Foo"),
+        ]
+    )
+    chain = prompt | ChatOpenAI(
+        api_key=OPENAI_API_KEY,
+        model="gpt-4o-mini",
+        temperature=0,
+        max_tokens=1,
+    )
+    callbacks = CallbackHandler(
+        mock_client,
+        trace_id="test-trace-id",
+        distinct_id="test_id",
+        properties={"foo": "bar"},
+    )
+    start_time = time.time()
+    result = chain.invoke({}, config={"callbacks": [callbacks]})
+    approximate_latency = math.floor(time.time() - start_time)
+
+    assert result.content == "Bar"
+    assert mock_client.capture.call_count == 2
+
+    first_call_args = mock_client.capture.call_args_list[0][1]
+    first_call_props = first_call_args["properties"]
+    assert first_call_args["event"] == "$ai_generation"
+    assert first_call_props["$ai_trace_id"] == "test-trace-id"
+    assert first_call_props["$ai_provider"] == "openai"
+    assert first_call_props["$ai_model"] == "gpt-4o-mini"
+    assert first_call_props["foo"] == "bar"
+
+    # langchain-openai for langchain v3
+    if "max_completion_tokens" in first_call_props["$ai_model_parameters"]:
+        assert first_call_props["$ai_model_parameters"] == {
+            "temperature": 0.0,
+            "max_completion_tokens": 1,
+            "stream": False,
+        }
+    else:
+        assert first_call_props["$ai_model_parameters"] == {
+            "temperature": 0.0,
+            "max_tokens": 1,
+            "n": 1,
+            "stream": False,
+        }
+    assert first_call_props["$ai_input"] == [
+        {"role": "system", "content": 'You must always answer with "Bar".'},
+        {"role": "user", "content": "Foo"},
+    ]
+    assert first_call_props["$ai_output_choices"] == [{"role": "assistant", "content": "Bar", "refusal": None}]
+    assert first_call_props["$ai_http_status"] == 200
+    assert isinstance(first_call_props["$ai_latency"], float)
+    assert min(approximate_latency - 1, 0) <= math.floor(first_call_props["$ai_latency"]) <= approximate_latency
+    assert first_call_props["$ai_input_tokens"] == 20
+    assert first_call_props["$ai_output_tokens"] == 1
+
+
+@pytest.mark.skipif(not OPENAI_API_KEY, reason="OpenAI API key not set")
+def test_openai_captures_multiple_generations(mock_client):
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", 'You must always answer with "Bar".'),
+            ("user", "Foo"),
+        ]
+    )
+    chain = prompt | ChatOpenAI(
+        api_key=OPENAI_API_KEY,
+        model="gpt-4o-mini",
+        temperature=0,
+        max_tokens=1,
+        n=2,
+    )
+    callbacks = CallbackHandler(mock_client)
+    result = chain.invoke({}, config={"callbacks": [callbacks]})
+
+    assert result.content == "Bar"
+    assert mock_client.capture.call_count == 2
+
+    first_call_args = mock_client.capture.call_args_list[0][1]
+    first_call_props = first_call_args["properties"]
+    second_call_args = mock_client.capture.call_args_list[1][1]
+    second_call_props = second_call_args["properties"]
+
+    assert first_call_args["event"] == "$ai_generation"
+    assert first_call_props["$ai_input"] == [
+        {"role": "system", "content": 'You must always answer with "Bar".'},
+        {"role": "user", "content": "Foo"},
+    ]
+    assert first_call_props["$ai_output_choices"] == [
+        {"role": "assistant", "content": "Bar", "refusal": None},
+        {
+            "role": "assistant",
+            "content": "Bar",
+        },
+    ]
+
+    # langchain-openai for langchain v3
+    if "max_completion_tokens" in first_call_props["$ai_model_parameters"]:
+        assert first_call_props["$ai_model_parameters"] == {
+            "temperature": 0.0,
+            "max_completion_tokens": 1,
+            "stream": False,
+            "n": 2,
+        }
+    else:
+        assert first_call_props["$ai_model_parameters"] == {
+            "temperature": 0.0,
+            "max_tokens": 1,
+            "stream": False,
+            "n": 2,
+        }
+    assert first_call_props["$ai_http_status"] == 200
+
+    assert second_call_args["event"] == "$ai_trace"
+    assert second_call_props["$ai_input_state"] == {}
+    assert isinstance(second_call_props["$ai_output_state"], AIMessage)
+
+
+@pytest.mark.skipif(not OPENAI_API_KEY, reason="OpenAI API key not set")
+def test_openai_streaming(mock_client):
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", 'You must always answer with "Bar".'),
+            ("user", "Foo"),
+        ]
+    )
+    chain = prompt | ChatOpenAI(
+        api_key=OPENAI_API_KEY,
+        model="gpt-4o-mini",
+        temperature=0,
+        max_tokens=1,
+        stream=True,
+        stream_usage=True,
+    )
+    callbacks = CallbackHandler(mock_client)
+    result = [m for m in chain.stream({}, config={"callbacks": [callbacks]})]
+    result = sum(result[1:], result[0])
+
+    assert result.content == "Bar"
+    assert mock_client.capture.call_count == 2
+
+    first_call_args = mock_client.capture.call_args_list[0][1]
+    first_call_props = first_call_args["properties"]
+    second_call_args = mock_client.capture.call_args_list[1][1]
+    second_call_props = second_call_args["properties"]
+
+    assert first_call_args["event"] == "$ai_generation"
+    assert first_call_props["$ai_model_parameters"]["stream"]
+    assert first_call_props["$ai_input"] == [
+        {"role": "system", "content": 'You must always answer with "Bar".'},
+        {"role": "user", "content": "Foo"},
+    ]
+    assert first_call_props["$ai_output_choices"] == [{"role": "assistant", "content": "Bar"}]
+    assert first_call_props["$ai_http_status"] == 200
+    assert first_call_props["$ai_input_tokens"] == 20
+    assert first_call_props["$ai_output_tokens"] == 1
+
+    assert second_call_args["event"] == "$ai_trace"
+    assert second_call_props["$ai_input_state"] == {"input": ""}
+    assert isinstance(second_call_props["$ai_output_state"], AIMessage)
+
+
+@pytest.mark.skipif(not OPENAI_API_KEY, reason="OpenAI API key not set")
+async def test_async_openai_streaming(mock_client):
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", 'You must always answer with "Bar".'),
+            ("user", "Foo"),
+        ]
+    )
+    chain = prompt | ChatOpenAI(
+        api_key=OPENAI_API_KEY,
+        model="gpt-4o-mini",
+        temperature=0,
+        max_tokens=1,
+        stream=True,
+        stream_usage=True,
+    )
+    callbacks = CallbackHandler(mock_client)
+    result = [m async for m in chain.astream({}, config={"callbacks": [callbacks]})]
+    result = sum(result[1:], result[0])
+
+    assert result.content == "Bar"
+    assert mock_client.capture.call_count == 2
+
+    first_call_args = mock_client.capture.call_args_list[0][1]
+    first_call_props = first_call_args["properties"]
+    second_call_args = mock_client.capture.call_args_list[1][1]
+    second_call_props = second_call_args["properties"]
+
+    assert first_call_args["event"] == "$ai_generation"
+    assert first_call_props["$ai_model_parameters"]["stream"]
+    assert first_call_props["$ai_input"] == [
+        {"role": "system", "content": 'You must always answer with "Bar".'},
+        {"role": "user", "content": "Foo"},
+    ]
+    assert first_call_props["$ai_output_choices"] == [{"role": "assistant", "content": "Bar"}]
+    assert first_call_props["$ai_http_status"] == 200
+    assert first_call_props["$ai_input_tokens"] == 20
+    assert first_call_props["$ai_output_tokens"] == 1
+
+    assert second_call_args["event"] == "$ai_trace"
+    assert second_call_props["$ai_input_state"] == {"input": ""}
+    assert isinstance(second_call_props["$ai_output_state"], AIMessage)
+
+
+def test_base_url_retrieval(mock_client):
+    prompt = ChatPromptTemplate.from_messages([("user", "Foo")])
+    chain = prompt | ChatOpenAI(
+        api_key="test",
+        model="posthog-mini",
+        base_url="https://test.posthog.com",
+    )
+    callbacks = CallbackHandler(mock_client)
+    with pytest.raises(Exception):
+        chain.invoke({}, config={"callbacks": [callbacks]})
+
+    assert mock_client.capture.call_count == 2
+    generation_call = mock_client.capture.call_args_list[0][1]
+    assert generation_call["properties"]["$ai_base_url"] == "https://test.posthog.com"
+
+
+def test_groups(mock_client):
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", 'You must always answer with "Bar".'),
+            ("user", "Foo"),
+        ]
+    )
+    model = FakeMessagesListChatModel(responses=[AIMessage(content="Bar")])
+    chain = prompt | model
+    callbacks = CallbackHandler(mock_client, groups={"company": "test_company"})
+    chain.invoke({}, config={"callbacks": [callbacks]})
+
+    assert mock_client.capture.call_count == 2
+    generation_call = mock_client.capture.call_args_list[0][1]
+    assert generation_call["groups"] == {"company": "test_company"}
+
+
+def test_privacy_mode_local(mock_client):
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", 'You must always answer with "Bar".'),
+            ("user", "Foo"),
+        ]
+    )
+    model = FakeMessagesListChatModel(responses=[AIMessage(content="Bar")])
+    chain = prompt | model
+    callbacks = CallbackHandler(mock_client, privacy_mode=True)
+    chain.invoke({}, config={"callbacks": [callbacks]})
+
+    assert mock_client.capture.call_count == 2
+    generation_call = mock_client.capture.call_args_list[0][1]
+    assert generation_call["properties"]["$ai_input"] is None
+    assert generation_call["properties"]["$ai_output_choices"] is None
+
+
+def test_privacy_mode_global(mock_client):
+    mock_client.privacy_mode = True
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", 'You must always answer with "Bar".'),
+            ("user", "Foo"),
+        ]
+    )
+    model = FakeMessagesListChatModel(responses=[AIMessage(content="Bar")])
+    chain = prompt | model
+    callbacks = CallbackHandler(mock_client)
+    chain.invoke({}, config={"callbacks": [callbacks]})
+
+    assert mock_client.capture.call_count == 2
+    generation_call = mock_client.capture.call_args_list[0][1]
+    assert generation_call["properties"]["$ai_input"] is None
+    assert generation_call["properties"]["$ai_output_choices"] is None
+
+
+@pytest.mark.skipif(not ANTHROPIC_API_KEY, reason="ANTHROPIC_API_KEY is not set")
+def test_anthropic_chain(mock_client):
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", 'You must always answer with "Bar".'),
+            ("user", "Foo"),
+        ]
+    )
+    chain = prompt | ChatAnthropic(
+        api_key=ANTHROPIC_API_KEY,
+        model="claude-3-opus-20240229",
+        temperature=0,
+        max_tokens=1,
+    )
+    callbacks = CallbackHandler(
+        mock_client,
+        trace_id="test-trace-id",
+        distinct_id="test_id",
+        properties={"foo": "bar"},
+    )
+    start_time = time.time()
+    result = chain.invoke({}, config={"callbacks": [callbacks]})
+    approximate_latency = math.floor(time.time() - start_time)
+
+    assert result.content == "Bar"
+    assert mock_client.capture.call_count == 2
+
+    first_call_args = mock_client.capture.call_args_list[0][1]
+    first_call_props = first_call_args["properties"]
+    second_call_args = mock_client.capture.call_args_list[1][1]
+    second_call_props = second_call_args["properties"]
+
+    assert first_call_args["event"] == "$ai_generation"
+    assert first_call_props["$ai_trace_id"] == "test-trace-id"
+    assert first_call_props["$ai_provider"] == "anthropic"
+    assert first_call_props["$ai_model"] == "claude-3-opus-20240229"
+    assert first_call_props["foo"] == "bar"
+
+    assert first_call_props["$ai_model_parameters"] == {
+        "temperature": 0.0,
+        "max_tokens": 1,
+        "streaming": False,
+    }
+    assert first_call_props["$ai_input"] == [
+        {"role": "system", "content": 'You must always answer with "Bar".'},
+        {"role": "user", "content": "Foo"},
+    ]
+    assert first_call_props["$ai_output_choices"] == [{"role": "assistant", "content": "Bar"}]
+    assert first_call_props["$ai_http_status"] == 200
+    assert isinstance(first_call_props["$ai_latency"], float)
+    assert min(approximate_latency - 1, 0) <= math.floor(first_call_props["$ai_latency"]) <= approximate_latency
+    assert first_call_props["$ai_input_tokens"] == 17
+    assert first_call_props["$ai_output_tokens"] == 1
+
+    assert second_call_args["event"] == "$ai_trace"
+    assert second_call_props["$ai_input_state"] == {}
+    assert isinstance(second_call_props["$ai_output_state"], AIMessage)
+
+
+@pytest.mark.skipif(not ANTHROPIC_API_KEY, reason="ANTHROPIC_API_KEY is not set")
+async def test_async_anthropic_streaming(mock_client):
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", 'You must always answer with "Bar".'),
+            ("user", "Foo"),
+        ]
+    )
+    chain = prompt | ChatAnthropic(
+        api_key=ANTHROPIC_API_KEY,
+        model="claude-3-opus-20240229",
+        temperature=0,
+        max_tokens=1,
+        streaming=True,
+        stream_usage=True,
+    )
+    callbacks = CallbackHandler(mock_client)
+    result = [m async for m in chain.astream({}, config={"callbacks": [callbacks]})]
+    result = sum(result[1:], result[0])
+
+    assert result.content == "Bar"
+    assert mock_client.capture.call_count == 2
+
+    first_call_args = mock_client.capture.call_args_list[0][1]
+    first_call_props = first_call_args["properties"]
+    second_call_args = mock_client.capture.call_args_list[1][1]
+    second_call_props = second_call_args["properties"]
+
+    assert first_call_args["event"] == "$ai_generation"
+    assert first_call_props["$ai_model_parameters"]["streaming"]
+    assert first_call_props["$ai_input"] == [
+        {"role": "system", "content": 'You must always answer with "Bar".'},
+        {"role": "user", "content": "Foo"},
+    ]
+    assert first_call_props["$ai_output_choices"] == [{"role": "assistant", "content": "Bar"}]
+    assert first_call_props["$ai_http_status"] == 200
+    assert first_call_props["$ai_input_tokens"] == 17
+    assert first_call_props["$ai_output_tokens"] is not None
+
+    assert second_call_args["event"] == "$ai_trace"
+    assert second_call_props["$ai_input_state"] == {
+        "input": "",
+    }
+    assert isinstance(second_call_props["$ai_output_state"], AIMessage)
+
+
+def test_tool_calls(mock_client):
+    prompt = ChatPromptTemplate.from_messages([("user", "Foo")])
+    model = FakeMessagesListChatModel(
+        responses=[
+            AIMessage(
+                content="Bar",
+                additional_kwargs={
+                    "tool_calls": [
+                        {
+                            "type": "function",
+                            "id": "123",
+                            "function": {
+                                "name": "test",
+                                "args": '{"a": 1}',
+                            },
+                        }
+                    ]
+                },
+            )
+        ]
+    )
+    chain = prompt | model
+    callbacks = CallbackHandler(mock_client)
+    chain.invoke({}, config={"callbacks": [callbacks]})
+
+    assert mock_client.capture.call_count == 2
+    generation_call = mock_client.capture.call_args_list[0][1]
+    assert generation_call["properties"]["$ai_output_choices"][0]["tool_calls"] == [
+        {
+            "type": "function",
+            "id": "123",
+            "function": {
+                "name": "test",
+                "args": '{"a": 1}',
+            },
+        }
+    ]
+    assert "additional_kwargs" not in generation_call["properties"]["$ai_output_choices"][0]

posthog/test/ai/openai/test_openai.py

@@ -0,0 +1,175 @@
+import time
+from unittest.mock import patch
+
+import pytest
+from openai.types.chat import ChatCompletion, ChatCompletionMessage
+from openai.types.chat.chat_completion import Choice
+from openai.types.completion_usage import CompletionUsage
+from openai.types.create_embedding_response import CreateEmbeddingResponse, Usage
+from openai.types.embedding import Embedding
+
+from posthog.ai.openai import OpenAI
+
+
+@pytest.fixture
+def mock_client():
+    with patch("posthog.client.Client") as mock_client:
+        mock_client.privacy_mode = False
+        yield mock_client
+
+
+@pytest.fixture
+def mock_openai_response():
+    return ChatCompletion(
+        id="test",
+        model="gpt-4",
+        object="chat.completion",
+        created=int(time.time()),
+        choices=[
+            Choice(
+                finish_reason="stop",
+                index=0,
+                message=ChatCompletionMessage(
+                    content="Test response",
+                    role="assistant",
+                ),
+            )
+        ],
+        usage=CompletionUsage(
+            completion_tokens=10,
+            prompt_tokens=20,
+            total_tokens=30,
+        ),
+    )
+
+
+@pytest.fixture
+def mock_embedding_response():
+    return CreateEmbeddingResponse(
+        data=[
+            Embedding(
+                embedding=[0.1, 0.2, 0.3],
+                index=0,
+                object="embedding",
+            )
+        ],
+        model="text-embedding-3-small",
+        object="list",
+        usage=Usage(
+            prompt_tokens=10,
+            total_tokens=10,
+        ),
+    )
+
+
+def test_basic_completion(mock_client, mock_openai_response):
+    with patch("openai.resources.chat.completions.Completions.create", return_value=mock_openai_response):
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+        response = client.chat.completions.create(
+            model="gpt-4",
+            messages=[{"role": "user", "content": "Hello"}],
+            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-4"
+        assert props["$ai_input"] == [{"role": "user", "content": "Hello"}]
+        assert props["$ai_output_choices"] == [{"role": "assistant", "content": "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)
+
+
+def test_embeddings(mock_client, mock_embedding_response):
+    with patch("openai.resources.embeddings.Embeddings.create", return_value=mock_embedding_response):
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+        response = client.embeddings.create(
+            model="text-embedding-3-small",
+            input="Hello world",
+            posthog_distinct_id="test-id",
+            posthog_properties={"foo": "bar"},
+        )
+
+        assert response == mock_embedding_response
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+
+        assert call_args["distinct_id"] == "test-id"
+        assert call_args["event"] == "$ai_embedding"
+        assert props["$ai_provider"] == "openai"
+        assert props["$ai_model"] == "text-embedding-3-small"
+        assert props["$ai_input"] == "Hello world"
+        assert props["$ai_input_tokens"] == 10
+        assert props["$ai_http_status"] == 200
+        assert props["foo"] == "bar"
+        assert isinstance(props["$ai_latency"], float)
+
+
+def test_groups(mock_client, mock_openai_response):
+    with patch("openai.resources.chat.completions.Completions.create", return_value=mock_openai_response):
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+        response = client.chat.completions.create(
+            model="gpt-4",
+            messages=[{"role": "user", "content": "Hello"}],
+            posthog_distinct_id="test-id",
+            posthog_groups={"company": "test_company"},
+        )
+
+        assert response == mock_openai_response
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+
+        assert call_args["groups"] == {"company": "test_company"}
+
+
+def test_privacy_mode_local(mock_client, mock_openai_response):
+    with patch("openai.resources.chat.completions.Completions.create", return_value=mock_openai_response):
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+        response = client.chat.completions.create(
+            model="gpt-4",
+            messages=[{"role": "user", "content": "Hello"}],
+            posthog_distinct_id="test-id",
+            posthog_privacy_mode=True,
+        )
+
+        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 props["$ai_input"] is None
+        assert props["$ai_output_choices"] is None
+
+
+def test_privacy_mode_global(mock_client, mock_openai_response):
+    with patch("openai.resources.chat.completions.Completions.create", return_value=mock_openai_response):
+        mock_client.privacy_mode = True
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+        response = client.chat.completions.create(
+            model="gpt-4",
+            messages=[{"role": "user", "content": "Hello"}],
+            posthog_distinct_id="test-id",
+            posthog_privacy_mode=False,
+        )
+
+        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 props["$ai_input"] is None
+        assert props["$ai_output_choices"] is None

posthog/test/test_client.py

@@ -104,11 +104,11 @@ class TestClient(unittest.TestCase):
         with mock.patch.object(Client, "capture", return_value=None) as patch_capture:
             client = self.client
             exception = Exception("test exception")
-            client.capture_exception(exception)
+            client.capture_exception(exception, distinct_id="distinct_id")
 
             self.assertTrue(patch_capture.called)
             capture_call = patch_capture.call_args[0]
-            self.assertEqual(capture_call[0], "python-exceptions")
+            self.assertEqual(capture_call[0], "distinct_id")
             self.assertEqual(capture_call[1], "$exception")
             self.assertEqual(
                 capture_call[2],
@@ -123,7 +123,7 @@ class TestClient(unittest.TestCase):
                             "value": "test exception",
                         }
                     ],
-                    "$exception_personURL": "https://us.i.posthog.com/project/random_key/person/python-exceptions",
+                    "$exception_personURL": "https://us.i.posthog.com/project/random_key/person/distinct_id",
                 },
             )
 
@@ -218,11 +218,11 @@ class TestClient(unittest.TestCase):
             try:
                 raise Exception("test exception")
             except Exception:
-                client.capture_exception()
+                client.capture_exception(distinct_id="distinct_id")
 
             self.assertTrue(patch_capture.called)
             capture_call = patch_capture.call_args[0]
-            self.assertEqual(capture_call[0], "python-exceptions")
+            self.assertEqual(capture_call[0], "distinct_id")
             self.assertEqual(capture_call[1], "$exception")
             self.assertEqual(capture_call[2]["$exception_type"], "Exception")
             self.assertEqual(capture_call[2]["$exception_message"], "test exception")
@@ -231,6 +231,10 @@ class TestClient(unittest.TestCase):
             self.assertEqual(capture_call[2]["$exception_list"][0]["module"], None)
             self.assertEqual(capture_call[2]["$exception_list"][0]["type"], "Exception")
             self.assertEqual(capture_call[2]["$exception_list"][0]["value"], "test exception")
+            self.assertEqual(
+                capture_call[2]["$exception_list"][0]["stacktrace"]["type"],
+                "raw",
+            )
             self.assertEqual(
                 capture_call[2]["$exception_list"][0]["stacktrace"]["frames"][0]["filename"],
                 "posthog/test/test_client.py",
@@ -711,6 +715,25 @@ class TestClient(unittest.TestCase):
         self.assertTrue(isinstance(msg["timestamp"], str))
         self.assertIsNone(msg.get("uuid"))
 
+    def test_basic_group_identify_with_distinct_id(self):
+        success, msg = self.client.group_identify("organization", "id:5", distinct_id="distinct_id")
+        self.assertTrue(success)
+        self.assertEqual(msg["event"], "$groupidentify")
+        self.assertEqual(msg["distinct_id"], "distinct_id")
+        self.assertEqual(
+            msg["properties"],
+            {
+                "$group_type": "organization",
+                "$group_key": "id:5",
+                "$group_set": {},
+                "$lib": "posthog-python",
+                "$lib_version": VERSION,
+                "$geoip_disable": True,
+            },
+        )
+        self.assertTrue(isinstance(msg["timestamp"], str))
+        self.assertIsNone(msg.get("uuid"))
+
     def test_advanced_group_identify(self):
         success, msg = self.client.group_identify(
             "organization", "id:5", {"trait": "value"}, {"ip": "192.168.0.1"}, datetime(2014, 9, 3), "new-uuid"
@@ -733,6 +756,35 @@ class TestClient(unittest.TestCase):
         self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00")
         self.assertEqual(msg["context"]["ip"], "192.168.0.1")
 
+    def test_advanced_group_identify_with_distinct_id(self):
+        success, msg = self.client.group_identify(
+            "organization",
+            "id:5",
+            {"trait": "value"},
+            {"ip": "192.168.0.1"},
+            datetime(2014, 9, 3),
+            "new-uuid",
+            distinct_id="distinct_id",
+        )
+
+        self.assertTrue(success)
+        self.assertEqual(msg["event"], "$groupidentify")
+        self.assertEqual(msg["distinct_id"], "distinct_id")
+
+        self.assertEqual(
+            msg["properties"],
+            {
+                "$group_type": "organization",
+                "$group_key": "id:5",
+                "$group_set": {"trait": "value"},
+                "$lib": "posthog-python",
+                "$lib_version": VERSION,
+                "$geoip_disable": True,
+            },
+        )
+        self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00")
+        self.assertEqual(msg["context"]["ip"], "192.168.0.1")
+
     def test_basic_alias(self):
         client = self.client
         success, msg = client.alias("previousId", "distinct_id")

posthog/test/test_exception_capture.py

@@ -29,7 +29,7 @@ def test_excepthook(tmpdir):
     assert b"LOL" in output
     assert b"DEBUG:posthog:data uploaded successfully" in output
     assert (
-        b'"$exception_list": [{"mechanism": {"type": "generic", "handled": true}, "module": null, "type": "ZeroDivisionError", "value": "division by zero", "stacktrace": {"frames": [{"filename": "app.py", "abs_path"'
+        b'"$exception_list": [{"mechanism": {"type": "generic", "handled": true}, "module": null, "type": "ZeroDivisionError", "value": "division by zero", "stacktrace": {"frames": [{"platform": "python", "filename": "app.py", "abs_path"'
         in output
     )
 
@@ -58,6 +58,6 @@ def test_trying_to_use_django_integration(tmpdir):
     assert b"LOL" in output
     assert b"DEBUG:posthog:data uploaded successfully" in output
     assert (
-        b'"$exception_list": [{"mechanism": {"type": "generic", "handled": true}, "module": null, "type": "ZeroDivisionError", "value": "division by zero", "stacktrace": {"frames": [{"filename": "app.py", "abs_path"'
+        b'"$exception_list": [{"mechanism": {"type": "generic", "handled": true}, "module": null, "type": "ZeroDivisionError", "value": "division by zero", "stacktrace": {"frames": [{"platform": "python", "filename": "app.py", "abs_path"'
         in output
     )

posthog/test/test_feature_flags.py

@@ -1632,9 +1632,10 @@ class TestLocalEvaluation(unittest.TestCase):
         )
         self.assertEqual(patch_decide.call_count, 0)
 
+    @mock.patch.object(Client, "capture")
     @mock.patch("posthog.client.decide")
-    def test_boolean_feature_flag_payload_decide(self, patch_decide):
-        patch_decide.return_value = {"featureFlagPayloads": {"person-flag": 300}}
+    def test_boolean_feature_flag_payload_decide(self, patch_decide, patch_capture):
+        patch_decide.return_value = {"featureFlags": {"person-flag": True}, "featureFlagPayloads": {"person-flag": 300}}
         self.assertEqual(
             self.client.get_feature_flag_payload(
                 "person-flag", "some-distinct-id", person_properties={"region": "USA"}
@@ -1649,6 +1650,8 @@ class TestLocalEvaluation(unittest.TestCase):
             300,
         )
         self.assertEqual(patch_decide.call_count, 2)
+        self.assertEqual(patch_capture.call_count, 1)
+        patch_capture.reset_mock()
 
     @mock.patch("posthog.client.decide")
     def test_multivariate_feature_flag_payloads(self, patch_decide):
@@ -2334,6 +2337,88 @@ class TestCaptureCalls(unittest.TestCase):
             disable_geoip=None,
         )
 
+    @mock.patch.object(Client, "capture")
+    @mock.patch("posthog.client.decide")
+    def test_capture_is_called_in_get_feature_flag_payload(self, patch_decide, patch_capture):
+        patch_decide.return_value = {
+            "featureFlags": {"person-flag": True},
+            "featureFlagPayloads": {"person-flag": 300},
+        }
+        client = Client(api_key=FAKE_TEST_API_KEY, personal_api_key=FAKE_TEST_API_KEY)
+
+        client.feature_flags = [
+            {
+                "id": 1,
+                "name": "Beta Feature",
+                "key": "person-flag",
+                "is_simple_flag": False,
+                "active": True,
+                "filters": {
+                    "groups": [
+                        {
+                            "properties": [{"key": "region", "value": "USA"}],
+                            "rollout_percentage": 100,
+                        }
+                    ],
+                },
+            }
+        ]
+
+        # Call get_feature_flag_payload with match_value=None to trigger get_feature_flag
+        client.get_feature_flag_payload(
+            key="person-flag", distinct_id="some-distinct-id", person_properties={"region": "USA", "name": "Aloha"}
+        )
+
+        # Assert that capture was called once, with the correct parameters
+        self.assertEqual(patch_capture.call_count, 1)
+        patch_capture.assert_called_with(
+            "some-distinct-id",
+            "$feature_flag_called",
+            {
+                "$feature_flag": "person-flag",
+                "$feature_flag_response": True,
+                "$feature_flag_payload": 300,
+                "locally_evaluated": False,
+                "$feature/person-flag": True,
+            },
+            groups={},
+            disable_geoip=None,
+        )
+
+        # Reset mocks for further tests
+        patch_capture.reset_mock()
+        patch_decide.reset_mock()
+
+        # Call get_feature_flag_payload again for the same user; capture should not be called again because we've already reported an event for this distinct_id + flag
+        client.get_feature_flag_payload(
+            key="person-flag", distinct_id="some-distinct-id", person_properties={"region": "USA", "name": "Aloha"}
+        )
+
+        self.assertEqual(patch_capture.call_count, 0)
+        patch_capture.reset_mock()
+
+        # Call get_feature_flag_payload for a different user; capture should be called
+        client.get_feature_flag_payload(
+            key="person-flag", distinct_id="some-distinct-id2", person_properties={"region": "USA", "name": "Aloha"}
+        )
+
+        self.assertEqual(patch_capture.call_count, 1)
+        patch_capture.assert_called_with(
+            "some-distinct-id2",
+            "$feature_flag_called",
+            {
+                "$feature_flag": "person-flag",
+                "$feature_flag_response": True,
+                "$feature_flag_payload": 300,
+                "locally_evaluated": False,
+                "$feature/person-flag": True,
+            },
+            groups={},
+            disable_geoip=None,
+        )
+
+        patch_capture.reset_mock()
+
     @mock.patch.object(Client, "capture")
     @mock.patch("posthog.client.decide")
     def test_disable_geoip_get_flag_capture_call(self, patch_decide, patch_capture):

posthog/version.py

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

setup.cfg

@@ -1,2 +1,5 @@
 [bdist_wheel]
 universal = 1
+
+[tool:pytest]
+asyncio_mode = auto

setup.py

@@ -14,7 +14,13 @@ long_description = """
 PostHog is developer-friendly, self-hosted product analytics. posthog-python is the python package.
 """
 
-install_requires = ["requests>=2.7,<3.0", "six>=1.5", "monotonic>=1.5", "backoff>=1.10.0", "python-dateutil>2.1"]
+install_requires = [
+    "requests>=2.7,<3.0",
+    "six>=1.5",
+    "monotonic>=1.5",
+    "backoff>=1.10.0",
+    "python-dateutil>2.1",
+]
 
 extras_require = {
     "dev": [
@@ -24,8 +30,25 @@ extras_require = {
         "flake8-print",
         "pre-commit",
     ],
-    "test": ["mock>=2.0.0", "freezegun==0.3.15", "pylint", "flake8", "coverage", "pytest", "pytest-timeout", "django"],
+    "test": [
+        "mock>=2.0.0",
+        "freezegun==0.3.15",
+        "pylint",
+        "flake8",
+        "coverage",
+        "pytest",
+        "pytest-timeout",
+        "pytest-asyncio",
+        "django",
+        "openai",
+        "anthropic",
+        "langgraph",
+        "langchain-community>=0.2.0",
+        "langchain-openai>=0.2.0",
+        "langchain-anthropic>=0.2.0",
+    ],
     "sentry": ["sentry-sdk", "django"],
+    "langchain": ["langchain>=0.2.0"],
 }
 
 setup(
@@ -37,7 +60,16 @@ setup(
     maintainer="PostHog",
     maintainer_email="hey@posthog.com",
     test_suite="posthog.test.all",
-    packages=["posthog", "posthog.test", "posthog.sentry", "posthog.exception_integrations"],
+    packages=[
+        "posthog",
+        "posthog.ai",
+        "posthog.ai.langchain",
+        "posthog.ai.openai",
+        "posthog.ai.anthropic",
+        "posthog.test",
+        "posthog.sentry",
+        "posthog.exception_integrations",
+    ],
     license="MIT License",
     install_requires=install_requires,
     extras_require=extras_require,
@@ -60,5 +92,8 @@ setup(
         "Programming Language :: Python :: 3.6",
         "Programming Language :: Python :: 3.7",
         "Programming Language :: Python :: 3.8",
+        "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
     ],
 )

setup_analytics.py

@@ -29,6 +29,10 @@ setup(
     test_suite="posthoganalytics.test.all",
     packages=[
         "posthoganalytics",
+        "posthoganalytics.ai",
+        "posthoganalytics.ai.langchain",
+        "posthoganalytics.ai.openai",
+        "posthoganalytics.ai.anthropic",
         "posthoganalytics.test",
         "posthoganalytics.sentry",
         "posthoganalytics.exception_integrations",
@@ -58,5 +62,8 @@ setup(
         "Programming Language :: Python :: 3.6",
         "Programming Language :: Python :: 3.7",
         "Programming Language :: Python :: 3.8",
+        "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
     ],
 )