fix: Actually fix LangChain callback in posthoganalytics

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

.github/workflows/ci.yml

@@ -36,7 +36,7 @@ jobs:
 
             - name: Lint with flake8
               run: |
-                  flake8 posthog --ignore E501
+                  flake8 posthog --ignore E501,W503
 
             - name: Check import order with isort
               run: |

CHANGELOG.md

@@ -1,4 +1,25 @@
-## 3.8.0 - 2025-01-14
+## 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.
 

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

llm_observability_examples.py

@@ -9,6 +9,8 @@ posthog.project_api_key = os.getenv("POSTHOG_PROJECT_API_KEY", "your-project-api
 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"),
@@ -26,11 +28,12 @@ def main_sync():
     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)
-        streaming_openai_call(distinct_id, trace_id, properties)
-        embedding_openai_call(distinct_id, trace_id, properties)
+        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))
@@ -41,17 +44,18 @@ async def main_async():
     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)
-        await streaming_async_openai_call(distinct_id, trace_id, properties)
-        await embedding_async_openai_call(distinct_id, trace_id, properties)
+        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):
+def basic_openai_call(distinct_id, trace_id, properties, groups):
     response = openai_client.chat.completions.create(
         model="gpt-4o-mini",
         messages=[
@@ -63,6 +67,7 @@ def basic_openai_call(distinct_id, trace_id, properties):
         posthog_distinct_id=distinct_id,
         posthog_trace_id=trace_id,
         posthog_properties=properties,
+        posthog_groups=groups,
     )
     print(response)
     if response and response.choices:
@@ -72,7 +77,7 @@ def basic_openai_call(distinct_id, trace_id, properties):
     return response
 
 
-async def basic_async_openai_call(distinct_id, trace_id, properties):
+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=[
@@ -84,6 +89,7 @@ async def basic_async_openai_call(distinct_id, trace_id, properties):
         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)
@@ -92,7 +98,7 @@ async def basic_async_openai_call(distinct_id, trace_id, properties):
     return response
 
 
-def streaming_openai_call(distinct_id, trace_id, properties):
+def streaming_openai_call(distinct_id, trace_id, properties, groups):
 
     response = openai_client.chat.completions.create(
         model="gpt-4o-mini",
@@ -106,6 +112,7 @@ def streaming_openai_call(distinct_id, trace_id, properties):
         posthog_distinct_id=distinct_id,
         posthog_trace_id=trace_id,
         posthog_properties=properties,
+        posthog_groups=groups,
     )
 
     for chunk in response:
@@ -115,7 +122,7 @@ def streaming_openai_call(distinct_id, trace_id, properties):
     return response
 
 
-async def streaming_async_openai_call(distinct_id, trace_id, properties):
+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=[
@@ -128,6 +135,7 @@ async def streaming_async_openai_call(distinct_id, trace_id, properties):
         posthog_distinct_id=distinct_id,
         posthog_trace_id=trace_id,
         posthog_properties=properties,
+        posthog_groups=groups,
     )
 
     async for chunk in response:
@@ -153,25 +161,27 @@ async def image_async_openai_call():
     return response
 
 
-def embedding_openai_call(posthog_distinct_id, posthog_trace_id, posthog_properties):
+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):
+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

posthog/__init__.py

@@ -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]
 

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/callbacks.py

@@ -19,11 +19,13 @@ from typing import (
 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.ai.utils import get_model_params
+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")
@@ -44,19 +46,30 @@ RunStorage = Dict[UUID, RunMetadata]
 
 class CallbackHandler(BaseCallbackHandler):
     """
-    A callback handler for LangChain that sends events to PostHog LLM Observability.
+    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),
@@ -65,10 +78,13 @@ class CallbackHandler(BaseCallbackHandler):
 
     def __init__(
         self,
-        client: Client,
+        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:
@@ -76,11 +92,17 @@ class CallbackHandler(BaseCallbackHandler):
             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
+        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 = {}
 
@@ -91,9 +113,14 @@ class CallbackHandler(BaseCallbackHandler):
         *,
         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,
@@ -104,6 +131,7 @@ class CallbackHandler(BaseCallbackHandler):
         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)
@@ -117,32 +145,93 @@ class CallbackHandler(BaseCallbackHandler):
         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,
-        tags: Optional[List[str]] = 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,
-        tags: Optional[List[str]] = 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)
@@ -164,8 +253,8 @@ class CallbackHandler(BaseCallbackHandler):
             "$ai_provider": run.get("provider"),
             "$ai_model": run.get("model"),
             "$ai_model_parameters": run.get("model_params"),
-            "$ai_input": run.get("messages"),
-            "$ai_output": {"choices": output},
+            "$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,
@@ -180,27 +269,18 @@ class CallbackHandler(BaseCallbackHandler):
             distinct_id=self._distinct_id or trace_id,
             event="$ai_generation",
             properties=event_properties,
+            groups=self._groups,
         )
 
-    def on_chain_error(
-        self,
-        error: BaseException,
-        *,
-        run_id: UUID,
-        parent_run_id: Optional[UUID] = None,
-        **kwargs: Any,
-    ):
-        self._pop_parent_of_run(run_id)
-
     def on_llm_error(
         self,
         error: BaseException,
         *,
         run_id: UUID,
         parent_run_id: Optional[UUID] = None,
-        tags: Optional[List[str]] = 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)
@@ -212,7 +292,7 @@ class CallbackHandler(BaseCallbackHandler):
             "$ai_provider": run.get("provider"),
             "$ai_model": run.get("model"),
             "$ai_model_parameters": run.get("model_params"),
-            "$ai_input": run.get("messages"),
+            "$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,
@@ -225,8 +305,53 @@ class CallbackHandler(BaseCallbackHandler):
             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.
@@ -296,6 +421,65 @@ class CallbackHandler(BaseCallbackHandler):
             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."""
@@ -325,15 +509,15 @@ def _convert_message_to_dict(message: BaseMessage) -> Dict[str, Any]:
     else:
         message_dict = {"role": message.type, "content": str(message.content)}
 
-    if "name" in message.additional_kwargs:
-        message_dict["name"] = message.additional_kwargs["name"]
     if message.additional_kwargs:
-        message_dict["additional_kwargs"] = 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]]:
+def _parse_usage_model(
+    usage: Union[BaseModel, Dict],
+) -> Tuple[Union[int, None], Union[int, None]]:
     if isinstance(usage, BaseModel):
         usage = usage.__dict__
 

posthog/ai/openai/openai.py

@@ -8,7 +8,7 @@ try:
 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
+from posthog.ai.utils import call_llm_and_track_usage, get_model_params, with_privacy_mode
 from posthog.client import Client as PostHogClient
 
 
@@ -49,6 +49,8 @@ class WrappedCompletions(openai.resources.chat.completions.Completions):
         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:
@@ -59,14 +61,19 @@ class WrappedCompletions(openai.resources.chat.completions.Completions):
                 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,
@@ -77,6 +84,8 @@ class WrappedCompletions(openai.resources.chat.completions.Completions):
         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()
@@ -117,6 +126,8 @@ class WrappedCompletions(openai.resources.chat.completions.Completions):
                     posthog_distinct_id,
                     posthog_trace_id,
                     posthog_properties,
+                    posthog_privacy_mode,
+                    posthog_groups,
                     kwargs,
                     usage_stats,
                     latency,
@@ -130,6 +141,8 @@ class WrappedCompletions(openai.resources.chat.completions.Completions):
         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,
@@ -142,15 +155,12 @@ class WrappedCompletions(openai.resources.chat.completions.Completions):
             "$ai_provider": "openai",
             "$ai_model": kwargs.get("model"),
             "$ai_model_parameters": get_model_params(kwargs),
-            "$ai_input": kwargs.get("messages"),
-            "$ai_output": {
-                "choices": [
-                    {
-                        "content": output,
-                        "role": "assistant",
-                    }
-                ]
-            },
+            "$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),
@@ -168,6 +178,7 @@ class WrappedCompletions(openai.resources.chat.completions.Completions):
                 distinct_id=posthog_distinct_id or posthog_trace_id,
                 event="$ai_generation",
                 properties=event_properties,
+                groups=posthog_groups,
             )
 
 
@@ -179,6 +190,8 @@ class WrappedEmbeddings(openai.resources.embeddings.Embeddings):
         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,
     ):
         """
@@ -214,7 +227,7 @@ class WrappedEmbeddings(openai.resources.embeddings.Embeddings):
         event_properties = {
             "$ai_provider": "openai",
             "$ai_model": kwargs.get("model"),
-            "$ai_input": kwargs.get("input"),
+            "$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,
@@ -232,6 +245,7 @@ class WrappedEmbeddings(openai.resources.embeddings.Embeddings):
                 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

@@ -8,7 +8,7 @@ try:
 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
+from posthog.ai.utils import call_llm_and_track_usage_async, get_model_params, with_privacy_mode
 from posthog.client import Client as PostHogClient
 
 
@@ -48,6 +48,8 @@ class WrappedCompletions(openai.resources.chat.completions.AsyncCompletions):
         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:
@@ -59,12 +61,15 @@ class WrappedCompletions(openai.resources.chat.completions.AsyncCompletions):
                 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,
@@ -78,6 +83,8 @@ class WrappedCompletions(openai.resources.chat.completions.AsyncCompletions):
         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()
@@ -112,10 +119,12 @@ class WrappedCompletions(openai.resources.chat.completions.AsyncCompletions):
                 end_time = time.time()
                 latency = end_time - start_time
                 output = "".join(accumulated_content)
-                self._capture_streaming_event(
+                await self._capture_streaming_event(
                     posthog_distinct_id,
                     posthog_trace_id,
                     posthog_properties,
+                    posthog_privacy_mode,
+                    posthog_groups,
                     kwargs,
                     usage_stats,
                     latency,
@@ -124,11 +133,13 @@ class WrappedCompletions(openai.resources.chat.completions.AsyncCompletions):
 
         return async_generator()
 
-    def _capture_streaming_event(
+    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,
@@ -141,15 +152,12 @@ class WrappedCompletions(openai.resources.chat.completions.AsyncCompletions):
             "$ai_provider": "openai",
             "$ai_model": kwargs.get("model"),
             "$ai_model_parameters": get_model_params(kwargs),
-            "$ai_input": kwargs.get("messages"),
-            "$ai_output": {
-                "choices": [
-                    {
-                        "content": output,
-                        "role": "assistant",
-                    }
-                ]
-            },
+            "$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),
@@ -167,6 +175,7 @@ class WrappedCompletions(openai.resources.chat.completions.AsyncCompletions):
                 distinct_id=posthog_distinct_id or posthog_trace_id,
                 event="$ai_generation",
                 properties=event_properties,
+                groups=posthog_groups,
             )
 
 
@@ -178,6 +187,8 @@ class WrappedEmbeddings(openai.resources.embeddings.AsyncEmbeddings):
         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,
     ):
         """
@@ -187,6 +198,8 @@ class WrappedEmbeddings(openai.resources.embeddings.AsyncEmbeddings):
             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:
@@ -213,7 +226,7 @@ class WrappedEmbeddings(openai.resources.embeddings.AsyncEmbeddings):
         event_properties = {
             "$ai_provider": "openai",
             "$ai_model": kwargs.get("model"),
-            "$ai_input": kwargs.get("input"),
+            "$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,
@@ -231,6 +244,7 @@ class WrappedEmbeddings(openai.resources.embeddings.AsyncEmbeddings):
                 distinct_id=posthog_distinct_id or posthog_trace_id,
                 event="$ai_embedding",
                 properties=event_properties,
+                groups=posthog_groups,
             )
 
         return response

posthog/ai/utils.py

@@ -21,23 +21,63 @@ def get_model_params(kwargs: Dict[str, Any]) -> Dict[str, Any]:
         "presence_penalty",
         "n",
         "stop",
-        "stream",
+        "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 format_response(response):
+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 = {"choices": []}
+    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["choices"].append(
+            output.append(
                 {
                     "content": choice.message.content,
                     "role": choice.message.role,
@@ -46,11 +86,23 @@ def format_response(response):
     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,
@@ -78,19 +130,21 @@ def call_llm_and_track_usage(
             posthog_trace_id = uuid.uuid4()
 
         if response and hasattr(response, "usage"):
-            usage = response.usage.model_dump()
+            usage = get_usage(response, provider)
+
+        messages = merge_system_prompt(kwargs, provider)
 
-        input_tokens = usage.get("prompt_tokens", 0)
-        output_tokens = usage.get("completion_tokens", 0)
         event_properties = {
-            "$ai_provider": "openai",
+            "$ai_provider": provider,
             "$ai_model": kwargs.get("model"),
             "$ai_model_parameters": get_model_params(kwargs),
-            "$ai_input": kwargs.get("messages"),
-            "$ai_output": format_response(response),
+            "$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": input_tokens,
-            "$ai_output_tokens": output_tokens,
+            "$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),
@@ -106,6 +160,7 @@ def call_llm_and_track_usage(
                 distinct_id=posthog_distinct_id or posthog_trace_id,
                 event="$ai_generation",
                 properties=event_properties,
+                groups=posthog_groups,
             )
 
     if error:
@@ -117,8 +172,11 @@ def call_llm_and_track_usage(
 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,
@@ -142,19 +200,21 @@ async def call_llm_and_track_usage_async(
             posthog_trace_id = uuid.uuid4()
 
         if response and hasattr(response, "usage"):
-            usage = response.usage.model_dump()
+            usage = get_usage(response, provider)
+
+        messages = merge_system_prompt(kwargs, provider)
 
-        input_tokens = usage.get("prompt_tokens", 0)
-        output_tokens = usage.get("completion_tokens", 0)
         event_properties = {
-            "$ai_provider": "openai",
+            "$ai_provider": provider,
             "$ai_model": kwargs.get("model"),
             "$ai_model_parameters": get_model_params(kwargs),
-            "$ai_input": kwargs.get("messages"),
-            "$ai_output": format_response(response),
+            "$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": input_tokens,
-            "$ai_output_tokens": output_tokens,
+            "$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),
@@ -170,9 +230,16 @@ async def call_llm_and_track_usage_async(
                 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

@@ -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:

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

@@ -2,3 +2,4 @@ import pytest
 
 pytest.importorskip("langchain")
 pytest.importorskip("langchain_community")
+pytest.importorskip("langgraph")

posthog/test/ai/langchain/test_callbacks.py

@@ -1,25 +1,32 @@
+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
+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
 
 
@@ -95,7 +102,11 @@ def test_basic_chat_chain(mock_client, stream):
         responses=[
             AIMessage(
                 content="The Los Angeles Dodgers won the World Series in 2020.",
-                usage_metadata={"input_tokens": 10, "output_tokens": 10, "total_tokens": 20},
+                usage_metadata={
+                    "input_tokens": 10,
+                    "output_tokens": 10,
+                    "total_tokens": 20,
+                },
             )
         ]
     )
@@ -107,26 +118,31 @@ def test_basic_chat_chain(mock_client, stream):
         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 == 1
-    args = mock_client.capture.call_args[1]
-    props = args["properties"]
+    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 args["event"] == "$ai_generation"
-    assert "distinct_id" in args
-    assert "$ai_model" in props
-    assert "$ai_provider" in props
-    assert props["$ai_input"] == [
+    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 props["$ai_output"] == {
-        "choices": [{"role": "assistant", "content": "The Los Angeles Dodgers won the World Series in 2020."}]
-    }
-    assert props["$ai_input_tokens"] == 10
-    assert props["$ai_output_tokens"] == 10
-    assert props["$ai_http_status"] == 200
-    assert props["$ai_trace_id"] is not None
-    assert isinstance(props["$ai_latency"], float)
+    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])
@@ -141,7 +157,11 @@ async def test_async_basic_chat_chain(mock_client, stream):
         responses=[
             AIMessage(
                 content="The Los Angeles Dodgers won the World Series in 2020.",
-                usage_metadata={"input_tokens": 10, "output_tokens": 10, "total_tokens": 20},
+                usage_metadata={
+                    "input_tokens": 10,
+                    "output_tokens": 10,
+                    "total_tokens": 20,
+                },
             )
         ]
     )
@@ -152,35 +172,50 @@ async def test_async_basic_chat_chain(mock_client, stream):
     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 == 1
+    assert mock_client.capture.call_count == 2
 
-    args = mock_client.capture.call_args[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"] == [
+    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 props["$ai_output"] == {
-        "choices": [{"role": "assistant", "content": "The Los Angeles Dodgers won the World Series in 2020."}]
-    }
-    assert props["$ai_input_tokens"] == 10
-    assert props["$ai_output_tokens"] == 10
-    assert props["$ai_http_status"] == 200
-    assert props["$ai_trace_id"] is not None
-    assert isinstance(props["$ai_latency"], float)
+    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)],
+    [
+        (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)]
+    callbacks: List[CallbackHandler] = [CallbackHandler(mock_client)]
 
     if stream:
         result = "".join(
@@ -191,7 +226,7 @@ def test_basic_llm_chain(mock_client, Model, stream):
     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[1]
+    args = mock_client.capture.call_args_list[0][1]
     props = args["properties"]
 
     assert args["event"] == "$ai_generation"
@@ -199,7 +234,7 @@ def test_basic_llm_chain(mock_client, Model, stream):
     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_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)
@@ -207,11 +242,16 @@ def test_basic_llm_chain(mock_client, Model, stream):
 
 @pytest.mark.parametrize(
     "Model,stream",
-    [(FakeListLLM, True), (FakeListLLM, False), (FakeStreamingListLLM, True), (FakeStreamingListLLM, False)],
+    [
+        (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)]
+    callbacks: List[CallbackHandler] = [CallbackHandler(mock_client)]
 
     if stream:
         result = "".join(
@@ -222,7 +262,7 @@ async def test_async_basic_llm_chain(mock_client, Model, stream):
     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[1]
+    args = mock_client.capture.call_args_list[0][1]
     props = args["properties"]
 
     assert args["event"] == "$ai_generation"
@@ -230,7 +270,7 @@ async def test_async_basic_llm_chain(mock_client, Model, stream):
     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_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)
@@ -248,7 +288,7 @@ def test_trace_id_for_multiple_chains(mock_client):
     result = chain.invoke({}, config={"callbacks": callbacks})
 
     assert result.content == "Bar"
-    assert mock_client.capture.call_count == 2
+    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"]
@@ -257,41 +297,59 @@ def test_trace_id_for_multiple_chains(mock_client):
     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_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_call_args = mock_client.capture.call_args_list[1][1]
-    second_call_props = second_call_args["properties"]
-    assert second_call_args["event"] == "$ai_generation"
-    assert "distinct_id" in second_call_args
-    assert "$ai_model" in second_call_props
-    assert "$ai_provider" in second_call_props
-    assert second_call_props["$ai_input"] == [{"role": "assistant", "content": "Bar"}]
-    assert second_call_props["$ai_output"] == {"choices": [{"role": "assistant", "content": "Bar"}]}
-    assert second_call_props["$ai_http_status"] == 200
-    assert second_call_props["$ai_trace_id"] is not None
-    assert isinstance(second_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_call_props["$ai_trace_id"]
+    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 == 1
-    args = mock_client.capture.call_args_list[0][1]
-    assert args["properties"]["$process_person_profile"] is False
+    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 == 2
-    args = mock_client.capture.call_args_list[1][1]
-    assert "$process_person_profile" not in args["properties"]
-    assert args["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):
@@ -300,17 +358,24 @@ def test_personless_mode_exception(mock_client):
     callbacks = CallbackHandler(mock_client)
     with pytest.raises(Exception):
         chain.invoke({}, config={"callbacks": [callbacks]})
-    assert mock_client.capture.call_count == 1
-    args = mock_client.capture.call_args_list[0][1]
-    assert args["properties"]["$process_person_profile"] is False
+    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 == 2
-    args = mock_client.capture.call_args_list[1][1]
-    assert "$process_person_profile" not in args["properties"]
-    assert args["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):
@@ -321,31 +386,127 @@ def test_metadata(mock_client):
     )
     model = FakeMessagesListChatModel(responses=[AIMessage(content="Bar")])
     callbacks = [
-        CallbackHandler(mock_client, trace_id="test-trace-id", distinct_id="test_id", properties={"foo": "bar"})
+        CallbackHandler(
+            mock_client,
+            trace_id="test-trace-id",
+            distinct_id="test_id",
+            properties={"foo": "bar"},
+        )
     ]
     chain = prompt | model
-    result = chain.invoke({}, config={"callbacks": callbacks})
+    result = chain.invoke({"plan": None}, config={"callbacks": callbacks})
 
     assert result.content == "Bar"
-    assert mock_client.capture.call_count == 1
+    assert mock_client.capture.call_count == 2
 
-    first_call_args = mock_client.capture.call_args[1]
-    assert first_call_args["distinct_id"] == "test_id"
+    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)
 
-    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["foo"] == "bar"
-    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 isinstance(first_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"})
+    callbacks = CallbackHandler(
+        mock_client,
+        trace_id="test-trace-id",
+        distinct_id="test_id",
+        properties={"foo": "bar"},
+    )
     chain = prompt | model
 
     chain.invoke({}, config={"callbacks": [callbacks]})
@@ -372,7 +533,10 @@ def test_exception_in_chain(mock_client):
 
     assert callbacks._runs == {}
     assert callbacks._parent_tree == {}
-    assert mock_client.capture.call_count == 0
+    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):
@@ -386,12 +550,12 @@ def test_openai_error(mock_client):
 
     assert callbacks._runs == {}
     assert callbacks._parent_tree == {}
-    assert mock_client.capture.call_count == 1
-    args = mock_client.capture.call_args[1]
-    props = args["properties"]
+    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" not in props
+    assert "$ai_output_choices" not in props
 
 
 @pytest.mark.skipif(not OPENAI_API_KEY, reason="OpenAI API key not set")
@@ -408,15 +572,20 @@ def test_openai_chain(mock_client):
         temperature=0,
         max_tokens=1,
     )
-    callbacks = CallbackHandler(mock_client, trace_id="test-trace-id", distinct_id="test_id", properties={"foo": "bar"})
+    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 == 1
+    assert mock_client.capture.call_count == 2
 
-    first_call_args = mock_client.capture.call_args[1]
+    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"
@@ -442,15 +611,7 @@ def test_openai_chain(mock_client):
         {"role": "system", "content": 'You must always answer with "Bar".'},
         {"role": "user", "content": "Foo"},
     ]
-    assert first_call_props["$ai_output"] == {
-        "choices": [
-            {
-                "role": "assistant",
-                "content": "Bar",
-                "additional_kwargs": {"refusal": None},
-            }
-        ]
-    }
+    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
@@ -477,27 +638,25 @@ def test_openai_captures_multiple_generations(mock_client):
     result = chain.invoke({}, config={"callbacks": [callbacks]})
 
     assert result.content == "Bar"
-    assert mock_client.capture.call_count == 1
+    assert mock_client.capture.call_count == 2
 
-    first_call_args = mock_client.capture.call_args[1]
+    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",
-                "additional_kwargs": {"refusal": None},
-            },
-            {
-                "role": "assistant",
-                "content": "Bar",
-            },
-        ]
-    }
+    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"]:
@@ -516,6 +675,10 @@ def test_openai_captures_multiple_generations(mock_client):
         }
     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):
@@ -526,28 +689,40 @@ def test_openai_streaming(mock_client):
         ]
     )
     chain = prompt | ChatOpenAI(
-        api_key=OPENAI_API_KEY, model="gpt-4o-mini", temperature=0, max_tokens=1, stream=True, stream_usage=True
+        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 == 1
+    assert mock_client.capture.call_count == 2
 
-    first_call_args = mock_client.capture.call_args[1]
+    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_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):
@@ -558,28 +733,40 @@ async def test_async_openai_streaming(mock_client):
         ]
     )
     chain = prompt | ChatOpenAI(
-        api_key=OPENAI_API_KEY, model="gpt-4o-mini", temperature=0, max_tokens=1, stream=True, stream_usage=True
+        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 == 1
+    assert mock_client.capture.call_count == 2
 
-    first_call_args = mock_client.capture.call_args[1]
+    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_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")])
@@ -592,6 +779,205 @@ def test_base_url_retrieval(mock_client):
     with pytest.raises(Exception):
         chain.invoke({}, config={"callbacks": [callbacks]})
 
-    assert mock_client.capture.call_count == 1
-    call = mock_client.capture.call_args[1]
-    assert call["properties"]["$ai_base_url"] == "https://test.posthog.com"
+    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

@@ -14,6 +14,7 @@ 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
 
 
@@ -82,7 +83,7 @@ def test_basic_completion(mock_client, mock_openai_response):
         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_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
@@ -115,3 +116,60 @@ def test_embeddings(mock_client, mock_embedding_response):
         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/version.py

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

setup.py

@@ -40,8 +40,12 @@ extras_require = {
         "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"],
@@ -58,6 +62,10 @@ setup(
     test_suite="posthog.test.all",
     packages=[
         "posthog",
+        "posthog.ai",
+        "posthog.ai.langchain",
+        "posthog.ai.openai",
+        "posthog.ai.anthropic",
         "posthog.test",
         "posthog.sentry",
         "posthog.exception_integrations",

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",
     ],
 )