Fix unit test

* feat: refactor to dataclasses

* feat: spans

* test: fix part 1

* test: fix part 2

* test: fix part 3

* test: fix part n

* test: add langgraph agent test

* chore: bump and linters

* chore: bump

* fix: correctly capture a parent id when a custom trace_id is set

* test: multiple spans parent_ids

* fix: exception serialization

* refactor: ai_trace_name -> ai_span_name and ai_generation_id -> ai_span_id

* fix: logs typos

* Add minor breaking change note to changelog

* fix: naming

---------

Co-authored-by: Michael Matloka <michael@matloka.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,3 +1,35 @@
+## 3.11.0 - 2025-01-28
+
+1. Add the `$ai_span` event to the LangChain callback handler to capture the input and output of intermediary chains.
+   > LLM observability naming change: event property `$ai_trace_name` is now `$ai_span_name`.
+
+2. Fix serialiazation of Pydantic models in methods.
+
+## 3.10.0 - 2025-01-24
+
+1. Add `$ai_error` and `$ai_is_error` properties to LangChain callback handler, OpenAI, and Anthropic.
+
+## 3.9.3 - 2025-01-23
+
+1. Fix capturing of multiple traces in the LangChain callback handler.
+
+## 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.

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]
 
@@ -34,7 +36,6 @@ def capture(
     distinct_id,  # type: str
     event,  # type: str
     properties=None,  # type: Optional[Dict]
-    context=None,  # type: Optional[Dict]
     timestamp=None,  # type: Optional[datetime.datetime]
     uuid=None,  # type: Optional[str]
     groups=None,  # type: Optional[Dict]
@@ -67,7 +68,6 @@ def capture(
         distinct_id=distinct_id,
         event=event,
         properties=properties,
-        context=context,
         timestamp=timestamp,
         uuid=uuid,
         groups=groups,
@@ -79,7 +79,6 @@ def capture(
 def identify(
     distinct_id,  # type: str
     properties=None,  # type: Optional[Dict]
-    context=None,  # type: Optional[Dict]
     timestamp=None,  # type: Optional[datetime.datetime]
     uuid=None,  # type: Optional[str]
     disable_geoip=None,  # type: Optional[bool]
@@ -104,7 +103,6 @@ def identify(
         "identify",
         distinct_id=distinct_id,
         properties=properties,
-        context=context,
         timestamp=timestamp,
         uuid=uuid,
         disable_geoip=disable_geoip,
@@ -114,7 +112,6 @@ def identify(
 def set(
     distinct_id,  # type: str
     properties=None,  # type: Optional[Dict]
-    context=None,  # type: Optional[Dict]
     timestamp=None,  # type: Optional[datetime.datetime]
     uuid=None,  # type: Optional[str]
     disable_geoip=None,  # type: Optional[bool]
@@ -139,7 +136,6 @@ def set(
         "set",
         distinct_id=distinct_id,
         properties=properties,
-        context=context,
         timestamp=timestamp,
         uuid=uuid,
         disable_geoip=disable_geoip,
@@ -149,7 +145,6 @@ def set(
 def set_once(
     distinct_id,  # type: str
     properties=None,  # type: Optional[Dict]
-    context=None,  # type: Optional[Dict]
     timestamp=None,  # type: Optional[datetime.datetime]
     uuid=None,  # type: Optional[str]
     disable_geoip=None,  # type: Optional[bool]
@@ -174,7 +169,6 @@ def set_once(
         "set_once",
         distinct_id=distinct_id,
         properties=properties,
-        context=context,
         timestamp=timestamp,
         uuid=uuid,
         disable_geoip=disable_geoip,
@@ -185,7 +179,6 @@ def group_identify(
     group_type,  # type: str
     group_key,  # type: str
     properties=None,  # type: Optional[Dict]
-    context=None,  # type: Optional[Dict]
     timestamp=None,  # type: Optional[datetime.datetime]
     uuid=None,  # type: Optional[str]
     disable_geoip=None,  # type: Optional[bool]
@@ -211,7 +204,6 @@ def group_identify(
         group_type=group_type,
         group_key=group_key,
         properties=properties,
-        context=context,
         timestamp=timestamp,
         uuid=uuid,
         disable_geoip=disable_geoip,
@@ -221,7 +213,6 @@ def group_identify(
 def alias(
     previous_id,  # type: str
     distinct_id,  # type: str
-    context=None,  # type: Optional[Dict]
     timestamp=None,  # type: Optional[datetime.datetime]
     uuid=None,  # type: Optional[str]
     disable_geoip=None,  # type: Optional[bool]
@@ -247,7 +238,6 @@ def alias(
         "alias",
         previous_id=previous_id,
         distinct_id=distinct_id,
-        context=context,
         timestamp=timestamp,
         uuid=uuid,
         disable_geoip=disable_geoip,

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

@@ -5,58 +5,93 @@ except ImportError:
 
 import logging
 import time
-import uuid
+from dataclasses import dataclass
 from typing import (
     Any,
     Dict,
     List,
     Optional,
+    Sequence,
     Tuple,
-    TypedDict,
     Union,
     cast,
 )
 from uuid import UUID
 
 from langchain.callbacks.base import BaseCallbackHandler
+from langchain.schema.agent import AgentAction, AgentFinish
+from langchain_core.documents import Document
 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")
 
 
-class RunMetadata(TypedDict, total=False):
-    messages: Union[List[Dict[str, Any]], List[str]]
-    provider: str
-    model: str
-    model_params: Dict[str, Any]
-    base_url: str
+@dataclass
+class SpanMetadata:
+    name: str
+    """Name of the run: chain name, model name, etc."""
     start_time: float
-    end_time: float
+    """Start time of the run."""
+    end_time: Optional[float]
+    """End time of the run."""
+    input: Optional[Any]
+    """Input of the run: messages, prompt variables, etc."""
+
+    @property
+    def latency(self) -> float:
+        if not self.end_time:
+            return 0
+        return self.end_time - self.start_time
 
 
-RunStorage = Dict[UUID, RunMetadata]
+@dataclass
+class GenerationMetadata(SpanMetadata):
+    provider: Optional[str] = None
+    """Provider of the run: OpenAI, Anthropic"""
+    model: Optional[str] = None
+    """Model used in the run"""
+    model_params: Optional[Dict[str, Any]] = None
+    """Model parameters of the run: temperature, max_tokens, etc."""
+    base_url: Optional[str] = None
+    """Base URL of the provider's API used in the run."""
+
+
+RunMetadata = Union[SpanMetadata, GenerationMetadata]
+RunMetadataStorage = 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
+
+    _runs: RunMetadataStorage
     """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 +100,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 +114,18 @@ 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
+        posthog_client = client or default_client
+        if posthog_client is None:
+            raise ValueError("PostHog client is required")
+        self._client = posthog_client
         self._distinct_id = distinct_id
         self._trace_id = trace_id
         self._properties = properties or {}
+        self._privacy_mode = privacy_mode
+        self._groups = groups or {}
         self._runs = {}
         self._parent_tree = {}
 
@@ -91,9 +136,34 @@ 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)
+        self._set_trace_or_span_metadata(serialized, inputs, run_id, parent_run_id, **kwargs)
+
+    def on_chain_end(
+        self,
+        outputs: Dict[str, Any],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ):
+        self._log_debug_event("on_chain_end", run_id, parent_run_id, outputs=outputs)
+        self._pop_run_and_capture_trace_or_span(run_id, parent_run_id, 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_run_and_capture_trace_or_span(run_id, parent_run_id, error)
 
     def on_chat_model_start(
         self,
@@ -104,9 +174,10 @@ 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)
+        self._set_llm_metadata(serialized, run_id, input, **kwargs)
 
     def on_llm_start(
         self,
@@ -117,19 +188,20 @@ 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)
+        self._set_llm_metadata(serialized, run_id, prompts, **kwargs)
 
-    def on_chain_end(
+    def on_llm_new_token(
         self,
-        outputs: Dict[str, Any],
+        token: str,
         *,
         run_id: UUID,
         parent_run_id: Optional[UUID] = None,
-        tags: Optional[List[str]] = None,
         **kwargs: Any,
-    ):
-        self._pop_parent_of_run(run_id)
+    ) -> 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_llm_end(
         self,
@@ -137,60 +209,13 @@ class CallbackHandler(BaseCallbackHandler):
         *,
         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.
         """
-        trace_id = self._get_trace_id(run_id)
-        self._pop_parent_of_run(run_id)
-        run = self._pop_run_metadata(run_id)
-        if not run:
-            return
-
-        latency = run.get("end_time", 0) - run.get("start_time", 0)
-        input_tokens, output_tokens = _parse_usage(response)
-
-        generation_result = response.generations[-1]
-        if isinstance(generation_result[-1], ChatGeneration):
-            output = [
-                _convert_message_to_dict(cast(ChatGeneration, generation).message) for generation in generation_result
-            ]
-        else:
-            output = [_extract_raw_esponse(generation) for generation in generation_result]
-
-        event_properties = {
-            "$ai_provider": run.get("provider"),
-            "$ai_model": run.get("model"),
-            "$ai_model_parameters": run.get("model_params"),
-            "$ai_input": run.get("messages"),
-            "$ai_output": {"choices": output},
-            "$ai_http_status": 200,
-            "$ai_input_tokens": input_tokens,
-            "$ai_output_tokens": output_tokens,
-            "$ai_latency": latency,
-            "$ai_trace_id": trace_id,
-            "$ai_base_url": run.get("base_url"),
-            **self._properties,
-        }
-        if self._distinct_id is None:
-            event_properties["$process_person_profile"] = False
-        self._client.capture(
-            distinct_id=self._distinct_id or trace_id,
-            event="$ai_generation",
-            properties=event_properties,
-        )
-
-    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)
+        self._log_debug_event("on_llm_end", run_id, parent_run_id, response=response, kwargs=kwargs)
+        self._pop_run_and_capture_generation(run_id, parent_run_id, response)
 
     def on_llm_error(
         self,
@@ -198,34 +223,107 @@ class CallbackHandler(BaseCallbackHandler):
         *,
         run_id: UUID,
         parent_run_id: Optional[UUID] = None,
-        tags: Optional[List[str]] = None,
         **kwargs: Any,
     ):
-        trace_id = self._get_trace_id(run_id)
-        self._pop_parent_of_run(run_id)
-        run = self._pop_run_metadata(run_id)
-        if not run:
-            return
+        self._log_debug_event("on_llm_error", run_id, parent_run_id, error=error)
+        self._pop_run_and_capture_generation(run_id, parent_run_id, error)
 
-        latency = run.get("end_time", 0) - run.get("start_time", 0)
-        event_properties = {
-            "$ai_provider": run.get("provider"),
-            "$ai_model": run.get("model"),
-            "$ai_model_parameters": run.get("model_params"),
-            "$ai_input": run.get("messages"),
-            "$ai_http_status": _get_http_status(error),
-            "$ai_latency": latency,
-            "$ai_trace_id": trace_id,
-            "$ai_base_url": run.get("base_url"),
-            **self._properties,
-        }
-        if self._distinct_id is None:
-            event_properties["$process_person_profile"] = False
-        self._client.capture(
-            distinct_id=self._distinct_id or trace_id,
-            event="$ai_generation",
-            properties=event_properties,
-        )
+    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)
+        self._set_trace_or_span_metadata(serialized, input_str, run_id, parent_run_id, **kwargs)
+
+    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)
+        self._pop_run_and_capture_trace_or_span(run_id, parent_run_id, output)
+
+    def on_tool_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[list[str]] = None,
+        **kwargs: Any,
+    ) -> Any:
+        self._log_debug_event("on_tool_error", run_id, parent_run_id, error=error)
+        self._pop_run_and_capture_trace_or_span(run_id, parent_run_id, error)
+
+    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)
+        self._set_trace_or_span_metadata(serialized, query, run_id, parent_run_id, **kwargs)
+
+    def on_retriever_end(
+        self,
+        documents: Sequence[Document],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ):
+        self._log_debug_event("on_retriever_end", run_id, parent_run_id, documents=documents)
+        self._pop_run_and_capture_trace_or_span(run_id, parent_run_id, documents)
+
+    def on_retriever_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[list[str]] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Run when Retriever errors."""
+        self._log_debug_event("on_retriever_error", run_id, parent_run_id, error=error)
+        self._pop_run_and_capture_trace_or_span(run_id, parent_run_id, 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)
+        self._set_parent_of_run(run_id, parent_run_id)
+        self._set_trace_or_span_metadata(None, action, run_id, parent_run_id, **kwargs)
+
+    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)
+        self._pop_run_and_capture_trace_or_span(run_id, parent_run_id, finish)
 
     def _set_parent_of_run(self, run_id: UUID, parent_run_id: Optional[UUID] = None):
         """
@@ -252,7 +350,19 @@ class CallbackHandler(BaseCallbackHandler):
             id = self._parent_tree[id]
         return id
 
-    def _set_run_metadata(
+    def _set_trace_or_span_metadata(
+        self,
+        serialized: Optional[Dict[str, Any]],
+        input: Any,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs,
+    ):
+        default_name = "trace" if parent_run_id is None else "span"
+        run_name = _get_langchain_run_name(serialized, **kwargs) or default_name
+        self._runs[run_id] = SpanMetadata(name=run_name, input=input, start_time=time.time(), end_time=None)
+
+    def _set_llm_metadata(
         self,
         serialized: Dict[str, Any],
         run_id: UUID,
@@ -261,24 +371,22 @@ class CallbackHandler(BaseCallbackHandler):
         invocation_params: Optional[Dict[str, Any]] = None,
         **kwargs,
     ):
-        run: RunMetadata = {
-            "messages": messages,
-            "start_time": time.time(),
-        }
+        run_name = _get_langchain_run_name(serialized, **kwargs) or "generation"
+        generation = GenerationMetadata(name=run_name, input=messages, start_time=time.time(), end_time=None)
         if isinstance(invocation_params, dict):
-            run["model_params"] = get_model_params(invocation_params)
+            generation.model_params = get_model_params(invocation_params)
         if isinstance(metadata, dict):
             if model := metadata.get("ls_model_name"):
-                run["model"] = model
+                generation.model = model
             if provider := metadata.get("ls_provider"):
-                run["provider"] = provider
+                generation.provider = provider
         try:
             base_url = serialized["kwargs"]["openai_api_base"]
             if base_url is not None:
-                run["base_url"] = base_url
+                generation.base_url = base_url
         except KeyError:
             pass
-        self._runs[run_id] = run
+        self._runs[run_id] = generation
 
     def _pop_run_metadata(self, run_id: UUID) -> Optional[RunMetadata]:
         end_time = time.time()
@@ -287,15 +395,155 @@ class CallbackHandler(BaseCallbackHandler):
         except KeyError:
             log.warning(f"No run metadata found for run {run_id}")
             return None
-        run["end_time"] = end_time
+        run.end_time = end_time
         return run
 
     def _get_trace_id(self, run_id: UUID):
         trace_id = self._trace_id or self._find_root_run(run_id)
         if not trace_id:
-            trace_id = uuid.uuid4()
+            return run_id
         return trace_id
 
+    def _get_parent_run_id(self, trace_id: Any, run_id: UUID, parent_run_id: Optional[UUID]):
+        """
+        Replace the parent run ID with the trace ID for second level runs when a custom trace ID is set.
+        """
+        if parent_run_id is not None and parent_run_id not in self._parent_tree:
+            return trace_id
+        return parent_run_id
+
+    def _pop_run_and_capture_trace_or_span(self, run_id: UUID, parent_run_id: Optional[UUID], outputs: Any):
+        trace_id = self._get_trace_id(run_id)
+        self._pop_parent_of_run(run_id)
+        run = self._pop_run_metadata(run_id)
+        if not run:
+            return
+        if isinstance(run, GenerationMetadata):
+            log.warning(f"Run {run_id} is a generation, but attempted to be captured as a trace or span.")
+            return
+        self._capture_trace_or_span(
+            trace_id, run_id, run, outputs, self._get_parent_run_id(trace_id, run_id, parent_run_id)
+        )
+
+    def _capture_trace_or_span(
+        self,
+        trace_id: Any,
+        run_id: UUID,
+        run: SpanMetadata,
+        outputs: Any,
+        parent_run_id: Optional[UUID],
+    ):
+        event_name = "$ai_trace" if parent_run_id is None else "$ai_span"
+        event_properties = {
+            "$ai_trace_id": trace_id,
+            "$ai_input_state": with_privacy_mode(self._client, self._privacy_mode, run.input),
+            "$ai_latency": run.latency,
+            "$ai_span_name": run.name,
+            "$ai_span_id": run_id,
+        }
+        if parent_run_id is not None:
+            event_properties["$ai_parent_id"] = parent_run_id
+        if self._properties:
+            event_properties.update(self._properties)
+
+        if isinstance(outputs, BaseException):
+            event_properties["$ai_error"] = _stringify_exception(outputs)
+            event_properties["$ai_is_error"] = True
+        elif outputs is not None:
+            event_properties["$ai_output_state"] = with_privacy_mode(self._client, self._privacy_mode, outputs)
+
+        if self._distinct_id is None:
+            event_properties["$process_person_profile"] = False
+
+        self._client.capture(
+            distinct_id=self._distinct_id or run_id,
+            event=event_name,
+            properties=event_properties,
+            groups=self._groups,
+        )
+
+    def _pop_run_and_capture_generation(
+        self, run_id: UUID, parent_run_id: Optional[UUID], response: Union[LLMResult, BaseException]
+    ):
+        trace_id = self._get_trace_id(run_id)
+        self._pop_parent_of_run(run_id)
+        run = self._pop_run_metadata(run_id)
+        if not run:
+            return
+        if not isinstance(run, GenerationMetadata):
+            log.warning(f"Run {run_id} is not a generation, but attempted to be captured as a generation.")
+            return
+        self._capture_generation(
+            trace_id, run_id, run, response, self._get_parent_run_id(trace_id, run_id, parent_run_id)
+        )
+
+    def _capture_generation(
+        self,
+        trace_id: Any,
+        run_id: UUID,
+        run: GenerationMetadata,
+        output: Union[LLMResult, BaseException],
+        parent_run_id: Optional[UUID] = None,
+    ):
+        event_properties = {
+            "$ai_trace_id": trace_id,
+            "$ai_span_id": run_id,
+            "$ai_span_name": run.name,
+            "$ai_parent_id": parent_run_id,
+            "$ai_provider": run.provider,
+            "$ai_model": run.model,
+            "$ai_model_parameters": run.model_params,
+            "$ai_input": with_privacy_mode(self._client, self._privacy_mode, run.input),
+            "$ai_http_status": 200,
+            "$ai_latency": run.latency,
+            "$ai_base_url": run.base_url,
+        }
+
+        if isinstance(output, BaseException):
+            event_properties["$ai_http_status"] = _get_http_status(output)
+            event_properties["$ai_error"] = _stringify_exception(output)
+            event_properties["$ai_is_error"] = True
+        else:
+            # Add usage
+            input_tokens, output_tokens = _parse_usage(output)
+            event_properties["$ai_input_tokens"] = input_tokens
+            event_properties["$ai_output_tokens"] = output_tokens
+
+            # Generation results
+            generation_result = output.generations[-1]
+            if isinstance(generation_result[-1], ChatGeneration):
+                completions = [
+                    _convert_message_to_dict(cast(ChatGeneration, generation).message)
+                    for generation in generation_result
+                ]
+            else:
+                completions = [_extract_raw_esponse(generation) for generation in generation_result]
+            event_properties["$ai_output_choices"] = with_privacy_mode(self._client, self._privacy_mode, completions)
+
+        if self._properties:
+            event_properties.update(self._properties)
+
+        if self._distinct_id is None:
+            event_properties["$process_person_profile"] = False
+
+        self._client.capture(
+            distinct_id=self._distinct_id or trace_id,
+            event="$ai_generation",
+            properties=event_properties,
+            groups=self._groups,
+        )
+
+    def _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 +573,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__
 
@@ -411,3 +659,41 @@ def _get_http_status(error: BaseException) -> int:
     # Google: https://github.com/googleapis/python-api-core/blob/main/google/api_core/exceptions.py
     status_code = getattr(error, "status_code", getattr(error, "code", 0))
     return status_code
+
+
+def _get_langchain_run_name(serialized: Optional[Dict[str, Any]], **kwargs: Any) -> Optional[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"]
+    if serialized is None:
+        return None
+    try:
+        return serialized["name"]
+    except (KeyError, TypeError):
+        pass
+    try:
+        return serialized["id"][-1]
+    except (KeyError, TypeError):
+        pass
+    return None
+
+
+def _stringify_exception(exception: BaseException) -> str:
+    description = str(exception)
+    if description:
+        return f"{exception.__class__.__name__}: {description}"
+    return exception.__class__.__name__

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,
@@ -64,12 +116,17 @@ def call_llm_and_track_usage(
     error = None
     http_status = 200
     usage: Dict[str, Any] = {}
+    error_params: Dict[str, any] = {}
 
     try:
         response = call_method(**kwargs)
     except Exception as exc:
         error = exc
         http_status = getattr(exc, "status_code", 0)  # default to 0 becuase its likely an SDK error
+        error_params = {
+            "$ai_is_error": True,
+            "$ai_error": exc.__str__(),
+        }
     finally:
         end_time = time.time()
         latency = end_time - start_time
@@ -78,23 +135,26 @@ 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),
             **(posthog_properties or {}),
+            **(error_params or {}),
         }
 
         if posthog_distinct_id is None:
@@ -106,6 +166,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 +178,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,
@@ -128,12 +192,17 @@ async def call_llm_and_track_usage_async(
     error = None
     http_status = 200
     usage: Dict[str, Any] = {}
+    error_params: Dict[str, any] = {}
 
     try:
         response = await call_async_method(**kwargs)
     except Exception as exc:
         error = exc
         http_status = getattr(exc, "status_code", 0)  # default to 0 because its likely an SDK error
+        error_params = {
+            "$ai_is_error": True,
+            "$ai_error": exc.__str__(),
+        }
     finally:
         end_time = time.time()
         latency = end_time - start_time
@@ -142,23 +211,26 @@ 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),
             **(posthog_properties or {}),
+            **(error_params or {}),
         }
 
         if posthog_distinct_id is None:
@@ -170,9 +242,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:
@@ -144,15 +146,13 @@ class Client(object):
                 if send:
                     consumer.start()
 
-    def identify(self, distinct_id=None, properties=None, context=None, timestamp=None, uuid=None, disable_geoip=None):
+    def identify(self, distinct_id=None, properties=None, timestamp=None, uuid=None, disable_geoip=None):
         properties = properties or {}
-        context = context or {}
         require("distinct_id", distinct_id, ID_TYPES)
         require("properties", properties, dict)
 
         msg = {
             "timestamp": timestamp,
-            "context": context,
             "distinct_id": distinct_id,
             "$set": properties,
             "event": "$identify",
@@ -209,7 +209,6 @@ class Client(object):
         distinct_id=None,
         event=None,
         properties=None,
-        context=None,
         timestamp=None,
         uuid=None,
         groups=None,
@@ -217,7 +216,6 @@ class Client(object):
         disable_geoip=None,
     ):
         properties = properties or {}
-        context = context or {}
         require("distinct_id", distinct_id, ID_TYPES)
         require("properties", properties, dict)
         require("event", event, string_types)
@@ -225,7 +223,6 @@ class Client(object):
         msg = {
             "properties": properties,
             "timestamp": timestamp,
-            "context": context,
             "distinct_id": distinct_id,
             "event": event,
             "uuid": uuid,
@@ -261,15 +258,13 @@ class Client(object):
 
         return self._enqueue(msg, disable_geoip)
 
-    def set(self, distinct_id=None, properties=None, context=None, timestamp=None, uuid=None, disable_geoip=None):
+    def set(self, distinct_id=None, properties=None, timestamp=None, uuid=None, disable_geoip=None):
         properties = properties or {}
-        context = context or {}
         require("distinct_id", distinct_id, ID_TYPES)
         require("properties", properties, dict)
 
         msg = {
             "timestamp": timestamp,
-            "context": context,
             "distinct_id": distinct_id,
             "$set": properties,
             "event": "$set",
@@ -278,15 +273,13 @@ class Client(object):
 
         return self._enqueue(msg, disable_geoip)
 
-    def set_once(self, distinct_id=None, properties=None, context=None, timestamp=None, uuid=None, disable_geoip=None):
+    def set_once(self, distinct_id=None, properties=None, timestamp=None, uuid=None, disable_geoip=None):
         properties = properties or {}
-        context = context or {}
         require("distinct_id", distinct_id, ID_TYPES)
         require("properties", properties, dict)
 
         msg = {
             "timestamp": timestamp,
-            "context": context,
             "distinct_id": distinct_id,
             "$set_once": properties,
             "event": "$set_once",
@@ -300,14 +293,12 @@ class Client(object):
         group_type=None,
         group_key=None,
         properties=None,
-        context=None,
         timestamp=None,
         uuid=None,
         disable_geoip=None,
         distinct_id=None,
     ):
         properties = properties or {}
-        context = context or {}
         require("group_type", group_type, ID_TYPES)
         require("group_key", group_key, ID_TYPES)
         require("properties", properties, dict)
@@ -326,15 +317,12 @@ class Client(object):
             },
             "distinct_id": distinct_id,
             "timestamp": timestamp,
-            "context": context,
             "uuid": uuid,
         }
 
         return self._enqueue(msg, disable_geoip)
 
-    def alias(self, previous_id=None, distinct_id=None, context=None, timestamp=None, uuid=None, disable_geoip=None):
-        context = context or {}
-
+    def alias(self, previous_id=None, distinct_id=None, timestamp=None, uuid=None, disable_geoip=None):
         require("previous_id", previous_id, ID_TYPES)
         require("distinct_id", distinct_id, ID_TYPES)
 
@@ -344,18 +332,14 @@ class Client(object):
                 "alias": distinct_id,
             },
             "timestamp": timestamp,
-            "context": context,
             "event": "$create_alias",
             "distinct_id": previous_id,
         }
 
         return self._enqueue(msg, disable_geoip)
 
-    def page(
-        self, distinct_id=None, url=None, properties=None, context=None, timestamp=None, uuid=None, disable_geoip=None
-    ):
+    def page(self, distinct_id=None, url=None, properties=None, timestamp=None, uuid=None, disable_geoip=None):
         properties = properties or {}
-        context = context or {}
 
         require("distinct_id", distinct_id, ID_TYPES)
         require("properties", properties, dict)
@@ -367,7 +351,6 @@ class Client(object):
             "event": "$pageview",
             "properties": properties,
             "timestamp": timestamp,
-            "context": context,
             "distinct_id": distinct_id,
             "uuid": uuid,
         }
@@ -444,7 +427,6 @@ class Client(object):
             timestamp = datetime.now(tz=tzutc())
 
         require("timestamp", timestamp, datetime)
-        require("context", msg["context"], dict)
 
         # add common
         timestamp = guess_timezone(timestamp)

posthog/test/ai/anthropic/test_anthropic.py

@@ -0,0 +1,341 @@
+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"},
+    ]
+
+
+def test_error(mock_client, mock_anthropic_response):
+    with patch("anthropic.resources.Messages.create", side_effect=Exception("Test error")):
+        client = Anthropic(api_key="test-key", posthog_client=mock_client)
+        with pytest.raises(Exception):
+            client.messages.create(model="claude-3-opus-20240229", messages=[{"role": "user", "content": "Hello"}])
+
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+        assert props["$ai_is_error"] is True
+        assert props["$ai_error"] == "Test error"

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/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,74 @@ 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
+
+
+def test_error(mock_client, mock_openai_response):
+    with patch("openai.resources.chat.completions.Completions.create", side_effect=Exception("Test error")):
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+        with pytest.raises(Exception):
+            client.chat.completions.create(model="gpt-4", messages=[{"role": "user", "content": "Hello"}])
+
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+        assert props["$ai_is_error"] is True
+        assert props["$ai_error"] == "Test error"

posthog/test/test_client.py

@@ -581,7 +581,6 @@ class TestClient(unittest.TestCase):
             "distinct_id",
             "python test event",
             {"property": "value"},
-            {"ip": "192.168.0.1"},
             datetime(2014, 9, 3),
             "new-uuid",
         )
@@ -590,7 +589,6 @@ class TestClient(unittest.TestCase):
 
         self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00")
         self.assertEqual(msg["properties"]["property"], "value")
-        self.assertEqual(msg["context"]["ip"], "192.168.0.1")
         self.assertEqual(msg["event"], "python test event")
         self.assertEqual(msg["properties"]["$lib"], "posthog-python")
         self.assertEqual(msg["properties"]["$lib_version"], VERSION)
@@ -622,14 +620,11 @@ class TestClient(unittest.TestCase):
 
     def test_advanced_identify(self):
         client = self.client
-        success, msg = client.identify(
-            "distinct_id", {"trait": "value"}, {"ip": "192.168.0.1"}, datetime(2014, 9, 3), "new-uuid"
-        )
+        success, msg = client.identify("distinct_id", {"trait": "value"}, datetime(2014, 9, 3), "new-uuid")
 
         self.assertTrue(success)
 
         self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00")
-        self.assertEqual(msg["context"]["ip"], "192.168.0.1")
         self.assertEqual(msg["$set"]["trait"], "value")
         self.assertEqual(msg["properties"]["$lib"], "posthog-python")
         self.assertEqual(msg["properties"]["$lib_version"], VERSION)
@@ -651,14 +646,11 @@ class TestClient(unittest.TestCase):
 
     def test_advanced_set(self):
         client = self.client
-        success, msg = client.set(
-            "distinct_id", {"trait": "value"}, {"ip": "192.168.0.1"}, datetime(2014, 9, 3), "new-uuid"
-        )
+        success, msg = client.set("distinct_id", {"trait": "value"}, datetime(2014, 9, 3), "new-uuid")
 
         self.assertTrue(success)
 
         self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00")
-        self.assertEqual(msg["context"]["ip"], "192.168.0.1")
         self.assertEqual(msg["$set"]["trait"], "value")
         self.assertEqual(msg["properties"]["$lib"], "posthog-python")
         self.assertEqual(msg["properties"]["$lib_version"], VERSION)
@@ -680,14 +672,11 @@ class TestClient(unittest.TestCase):
 
     def test_advanced_set_once(self):
         client = self.client
-        success, msg = client.set_once(
-            "distinct_id", {"trait": "value"}, {"ip": "192.168.0.1"}, datetime(2014, 9, 3), "new-uuid"
-        )
+        success, msg = client.set_once("distinct_id", {"trait": "value"}, datetime(2014, 9, 3), "new-uuid")
 
         self.assertTrue(success)
 
         self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00")
-        self.assertEqual(msg["context"]["ip"], "192.168.0.1")
         self.assertEqual(msg["$set_once"]["trait"], "value")
         self.assertEqual(msg["properties"]["$lib"], "posthog-python")
         self.assertEqual(msg["properties"]["$lib_version"], VERSION)
@@ -736,7 +725,7 @@ class TestClient(unittest.TestCase):
 
     def test_advanced_group_identify(self):
         success, msg = self.client.group_identify(
-            "organization", "id:5", {"trait": "value"}, {"ip": "192.168.0.1"}, datetime(2014, 9, 3), "new-uuid"
+            "organization", "id:5", {"trait": "value"}, datetime(2014, 9, 3), "new-uuid"
         )
 
         self.assertTrue(success)
@@ -754,14 +743,12 @@ class TestClient(unittest.TestCase):
             },
         )
         self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00")
-        self.assertEqual(msg["context"]["ip"], "192.168.0.1")
 
     def test_advanced_group_identify_with_distinct_id(self):
         success, msg = self.client.group_identify(
             "organization",
             "id:5",
             {"trait": "value"},
-            {"ip": "192.168.0.1"},
             datetime(2014, 9, 3),
             "new-uuid",
             distinct_id="distinct_id",
@@ -783,7 +770,6 @@ class TestClient(unittest.TestCase):
             },
         )
         self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00")
-        self.assertEqual(msg["context"]["ip"], "192.168.0.1")
 
     def test_basic_alias(self):
         client = self.client
@@ -819,7 +805,6 @@ class TestClient(unittest.TestCase):
             "distinct_id",
             "https://posthog.com/contact",
             {"property": "value"},
-            {"ip": "192.168.0.1"},
             datetime(2014, 9, 3),
             "new-uuid",
         )
@@ -827,7 +812,6 @@ class TestClient(unittest.TestCase):
         self.assertTrue(success)
 
         self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00")
-        self.assertEqual(msg["context"]["ip"], "192.168.0.1")
         self.assertEqual(msg["properties"]["$current_url"], "https://posthog.com/contact")
         self.assertEqual(msg["properties"]["property"], "value")
         self.assertEqual(msg["properties"]["$lib"], "posthog-python")

posthog/test/test_utils.py

@@ -1,10 +1,13 @@
 import unittest
 from datetime import date, datetime, timedelta
 from decimal import Decimal
+from typing import Optional
 from uuid import UUID
 
 import six
 from dateutil.tz import tzutc
+from pydantic import BaseModel
+from pydantic.v1 import BaseModel as BaseModelV1
 
 from posthog import utils
 
@@ -81,6 +84,32 @@ class TestUtils(unittest.TestCase):
         self.assertEqual("http://posthog.io", utils.remove_trailing_slash("http://posthog.io/"))
         self.assertEqual("http://posthog.io", utils.remove_trailing_slash("http://posthog.io"))
 
+    def test_clean_pydantic(self):
+        class ModelV2(BaseModel):
+            foo: str
+            bar: int
+            baz: Optional[str] = None
+
+        class ModelV1(BaseModelV1):
+            foo: int
+            bar: str
+
+        class NestedModel(BaseModel):
+            foo: ModelV2
+
+        self.assertEqual(utils.clean(ModelV2(foo="1", bar=2)), {"foo": "1", "bar": 2, "baz": None})
+        self.assertEqual(utils.clean(ModelV1(foo=1, bar="2")), {"foo": 1, "bar": "2"})
+        self.assertEqual(
+            utils.clean(NestedModel(foo=ModelV2(foo="1", bar=2, baz="3"))), {"foo": {"foo": "1", "bar": 2, "baz": "3"}}
+        )
+
+        class Dummy:
+            def model_dump(self, required_param):
+                pass
+
+        # Skips a class with a defined non-Pydantic `model_dump` method.
+        self.assertEqual(utils.clean({"test": Dummy()}), {})
+
 
 class TestSizeLimitedDict(unittest.TestCase):
     def test_size_limited_dict(self):

posthog/utils.py

@@ -51,14 +51,24 @@ def clean(item):
         return float(item)
     if isinstance(item, UUID):
         return str(item)
-    elif isinstance(item, (six.string_types, bool, numbers.Number, datetime, date, type(None))):
+    if isinstance(item, (six.string_types, bool, numbers.Number, datetime, date, type(None))):
         return item
-    elif isinstance(item, (set, list, tuple)):
+    if isinstance(item, (set, list, tuple)):
         return _clean_list(item)
-    elif isinstance(item, dict):
+    # Pydantic model
+    try:
+        # v2+
+        if hasattr(item, "model_dump") and callable(item.model_dump):
+            item = item.model_dump()
+        # v1
+        elif hasattr(item, "dict") and callable(item.dict):
+            item = item.dict()
+    except TypeError as e:
+        log.debug(f"Could not serialize Pydantic-like model: {e}")
+        pass
+    if isinstance(item, dict):
         return _clean_dict(item)
-    else:
-        return _coerce_unicode(item)
+    return _coerce_unicode(item)
 
 
 def _clean_list(list_):

posthog/version.py

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

setup.py

@@ -40,8 +40,13 @@ 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",
+        "pydantic",
     ],
     "sentry": ["sentry-sdk", "django"],
     "langchain": ["langchain>=0.2.0"],
@@ -61,6 +66,7 @@ setup(
         "posthog.ai",
         "posthog.ai.langchain",
         "posthog.ai.openai",
+        "posthog.ai.anthropic",
         "posthog.test",
         "posthog.sentry",
         "posthog.exception_integrations",

setup_analytics.py

@@ -30,6 +30,9 @@ setup(
     packages=[
         "posthoganalytics",
         "posthoganalytics.ai",
+        "posthoganalytics.ai.langchain",
+        "posthoganalytics.ai.openai",
+        "posthoganalytics.ai.anthropic",
         "posthoganalytics.test",
         "posthoganalytics.sentry",
         "posthoganalytics.exception_integrations",
@@ -59,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",
     ],
 )