fix(llm-observability): include the ai packages (#162 )

* fix: setuptools * fix: include packages
fix: update to export module (#161 )
2025-01-14 10:27:05 +01:00 · 2025-01-14 01:25:00 +00:00 · 2025-01-14 00:35:09 +00:00 · 2025-01-13 18:40:02 +01:00 · 2025-01-11 01:34:27 +00:00 · 2025-01-03 16:00:35 -05:00
23 changed files with 2091 additions and 14 deletions
@@ -18,7 +18,7 @@ jobs:
              with:
                  python-version: 3.8

-            - uses: actions/cache@v1
+            - uses: actions/cache@v3
              with:
                  path: ~/.cache/pip
                  key: ${{ runner.os }}-pip-${{ hashFiles('setup.py') }}
@@ -33,7 +33,7 @@ jobs:
            - name: Check formatting with black
              run: |
                  black --check .
-            
+
            - name: Lint with flake8
              run: |
                  flake8 posthog --ignore E501
@@ -47,14 +47,14 @@ jobs:
        runs-on: ubuntu-latest

        steps:
-            - uses: actions/checkout@v1
+            - uses: actions/checkout@v2
              with:
                  fetch-depth: 1

-            - name: Set up Python 3.7
-              uses: actions/setup-python@v1
+            - name: Set up Python 3.9
+              uses: actions/setup-python@v2
              with:
-                  python-version: 3.7
+                  python-version: 3.9

            - name: Install requirements.txt dependencies with pip
              run: |
@@ -14,4 +14,6 @@ pylint.out
 posthog-analytics
 .idea
 .python-version
-.coverage
+.coverage
+pyrightconfig.json
+.env
@@ -1,3 +1,15 @@
+## 3.8.2 - 2025-01-14
+
+1. Fix setuptools to include the `posthog.ai.openai` and `posthog.ai.langchain` packages.
+
+## 3.8.1 - 2025-01-14
+
+1. Add LLM Observability with support for OpenAI and Langchain callbacks.
+
+## 3.7.5 - 2025-01-03
+
+1. Add `distinct_id` to group_identify
+
 ## 3.7.4 - 2024-11-25

 1. Fix bug where this SDK incorrectly sent feature flag events with null values when calling `get_feature_flag_payload`.
@@ -1 +1 @@
-@PostHog/team-feature-success
+@PostHog/team-feature-flags
@@ -1,7 +1,7 @@
 # PostHog Python library example

 # Import the library
-import time
+# import time

 import posthog

@@ -0,0 +1,186 @@
+import os
+import uuid
+
+import posthog
+from posthog.ai.openai import AsyncOpenAI, OpenAI
+
+# Example credentials - replace these with your own or use environment variables
+posthog.project_api_key = os.getenv("POSTHOG_PROJECT_API_KEY", "your-project-api-key")
+posthog.personal_api_key = os.getenv("POSTHOG_PERSONAL_API_KEY", "your-personal-api-key")
+posthog.host = os.getenv("POSTHOG_HOST", "http://localhost:8000")  # Or https://app.posthog.com
+posthog.debug = True
+
+openai_client = OpenAI(
+    api_key=os.getenv("OPENAI_API_KEY", "your-openai-api-key"),
+    posthog_client=posthog,
+)
+
+async_openai_client = AsyncOpenAI(
+    api_key=os.getenv("OPENAI_API_KEY", "your-openai-api-key"),
+    posthog_client=posthog,
+)
+
+
+def main_sync():
+    trace_id = str(uuid.uuid4())
+    print("Trace ID:", trace_id)
+    distinct_id = "test2_distinct_id"
+    properties = {"test_property": "test_value"}
+
+    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)
+        image_openai_call()
+    except Exception as e:
+        print("Error during OpenAI call:", str(e))
+
+
+async def main_async():
+    trace_id = str(uuid.uuid4())
+    print("Trace ID:", trace_id)
+    distinct_id = "test_distinct_id"
+    properties = {"test_property": "test_value"}
+
+    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 image_async_openai_call()
+    except Exception as e:
+        print("Error during OpenAI call:", str(e))
+
+
+def basic_openai_call(distinct_id, trace_id, properties):
+    response = openai_client.chat.completions.create(
+        model="gpt-4o-mini",
+        messages=[
+            {"role": "system", "content": "You are a complex problem solver."},
+            {"role": "user", "content": "Explain quantum computing in simple terms."},
+        ],
+        max_tokens=100,
+        temperature=0.7,
+        posthog_distinct_id=distinct_id,
+        posthog_trace_id=trace_id,
+        posthog_properties=properties,
+    )
+    print(response)
+    if response and response.choices:
+        print("OpenAI response:", response.choices[0].message.content)
+    else:
+        print("No response or unexpected format returned.")
+    return response
+
+
+async def basic_async_openai_call(distinct_id, trace_id, properties):
+    response = await async_openai_client.chat.completions.create(
+        model="gpt-4o-mini",
+        messages=[
+            {"role": "system", "content": "You are a complex problem solver."},
+            {"role": "user", "content": "Explain quantum computing in simple terms."},
+        ],
+        max_tokens=100,
+        temperature=0.7,
+        posthog_distinct_id=distinct_id,
+        posthog_trace_id=trace_id,
+        posthog_properties=properties,
+    )
+    if response and hasattr(response, "choices"):
+        print("OpenAI response:", response.choices[0].message.content)
+    else:
+        print("No response or unexpected format returned.")
+    return response
+
+
+def streaming_openai_call(distinct_id, trace_id, properties):
+
+    response = openai_client.chat.completions.create(
+        model="gpt-4o-mini",
+        messages=[
+            {"role": "system", "content": "You are a complex problem solver."},
+            {"role": "user", "content": "Explain quantum computing in simple terms."},
+        ],
+        max_tokens=100,
+        temperature=0.7,
+        stream=True,
+        posthog_distinct_id=distinct_id,
+        posthog_trace_id=trace_id,
+        posthog_properties=properties,
+    )
+
+    for chunk in response:
+        if hasattr(chunk, "choices") and chunk.choices and len(chunk.choices) > 0:
+            print(chunk.choices[0].delta.content or "", end="")
+
+    return response
+
+
+async def streaming_async_openai_call(distinct_id, trace_id, properties):
+    response = await async_openai_client.chat.completions.create(
+        model="gpt-4o-mini",
+        messages=[
+            {"role": "system", "content": "You are a complex problem solver."},
+            {"role": "user", "content": "Explain quantum computing in simple terms."},
+        ],
+        max_tokens=100,
+        temperature=0.7,
+        stream=True,
+        posthog_distinct_id=distinct_id,
+        posthog_trace_id=trace_id,
+        posthog_properties=properties,
+    )
+
+    async for chunk in response:
+        if hasattr(chunk, "choices") and chunk.choices and len(chunk.choices) > 0:
+            print(chunk.choices[0].delta.content or "", end="")
+
+    return response
+
+
+# none instrumented
+def image_openai_call():
+    response = openai_client.images.generate(model="dall-e-3", prompt="A cute baby hedgehog", n=1, size="1024x1024")
+    print(response)
+    return response
+
+
+# none instrumented
+async def image_async_openai_call():
+    response = await async_openai_client.images.generate(
+        model="dall-e-3", prompt="A cute baby hedgehog", n=1, size="1024x1024"
+    )
+    print(response)
+    return response
+
+
+def embedding_openai_call(posthog_distinct_id, posthog_trace_id, posthog_properties):
+    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,
+    )
+    print(response)
+    return response
+
+
+async def embedding_async_openai_call(posthog_distinct_id, posthog_trace_id, posthog_properties):
+    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,
+    )
+    print(response)
+    return response
+
+
+# HOW TO RUN:
+# comment out one of these to run the other
+
+if __name__ == "__main__":
+    main_sync()
+
+# asyncio.run(main_async())
@@ -0,0 +1,3 @@
+from .callbacks import CallbackHandler
+
+__all__ = ["CallbackHandler"]
@@ -0,0 +1,413 @@
+try:
+    import langchain  # noqa: F401
+except ImportError:
+    raise ModuleNotFoundError("Please install LangChain to use this feature: 'pip install langchain'")
+
+import logging
+import time
+import uuid
+from typing import (
+    Any,
+    Dict,
+    List,
+    Optional,
+    Tuple,
+    TypedDict,
+    Union,
+    cast,
+)
+from uuid import UUID
+
+from langchain.callbacks.base import BaseCallbackHandler
+from langchain_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.client import Client
+
+log = logging.getLogger("posthog")
+
+
+class RunMetadata(TypedDict, total=False):
+    messages: Union[List[Dict[str, Any]], List[str]]
+    provider: str
+    model: str
+    model_params: Dict[str, Any]
+    base_url: str
+    start_time: float
+    end_time: float
+
+
+RunStorage = Dict[UUID, RunMetadata]
+
+
+class CallbackHandler(BaseCallbackHandler):
+    """
+    A callback handler for LangChain that sends events to PostHog LLM Observability.
+    """
+
+    _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."""
+    _properties: Optional[Dict[str, Any]]
+    """Global properties to be sent with every event."""
+    _runs: RunStorage
+    """Mapping of run IDs to run metadata as run metadata is only available on the start of generation."""
+    _parent_tree: Dict[UUID, UUID]
+    """
+    A dictionary that maps chain run IDs to their parent chain run IDs (parent pointer tree),
+    so the top level can be found from a bottom-level run ID.
+    """
+
+    def __init__(
+        self,
+        client: Client,
+        distinct_id: Optional[Union[str, int, float, UUID]] = None,
+        trace_id: Optional[Union[str, int, float, UUID]] = None,
+        properties: Optional[Dict[str, Any]] = None,
+    ):
+        """
+        Args:
+            client: PostHog client instance.
+            distinct_id: Optional distinct ID of the user to associate the trace with.
+            trace_id: Optional trace ID to use for the event.
+            properties: Optional additional metadata to use for the trace.
+        """
+        self._client = client
+        self._distinct_id = distinct_id
+        self._trace_id = trace_id
+        self._properties = properties or {}
+        self._runs = {}
+        self._parent_tree = {}
+
+    def on_chain_start(
+        self,
+        serialized: Dict[str, Any],
+        inputs: Dict[str, Any],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs,
+    ):
+        self._set_parent_of_run(run_id, parent_run_id)
+
+    def on_chat_model_start(
+        self,
+        serialized: Dict[str, Any],
+        messages: List[List[BaseMessage]],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs,
+    ):
+        self._set_parent_of_run(run_id, parent_run_id)
+        input = [_convert_message_to_dict(message) for row in messages for message in row]
+        self._set_run_metadata(serialized, run_id, input, **kwargs)
+
+    def on_llm_start(
+        self,
+        serialized: Dict[str, Any],
+        prompts: List[str],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ):
+        self._set_parent_of_run(run_id, parent_run_id)
+        self._set_run_metadata(serialized, run_id, prompts, **kwargs)
+
+    def on_chain_end(
+        self,
+        outputs: Dict[str, Any],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        **kwargs: Any,
+    ):
+        self._pop_parent_of_run(run_id)
+
+    def on_llm_end(
+        self,
+        response: LLMResult,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        **kwargs: Any,
+    ):
+        """
+        The callback works for both streaming and non-streaming runs. For streaming runs, the chain must set `stream_usage=True` in the LLM.
+        """
+        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)
+
+    def on_llm_error(
+        self,
+        error: BaseException,
+        *,
+        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
+
+        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 _set_parent_of_run(self, run_id: UUID, parent_run_id: Optional[UUID] = None):
+        """
+        Set the parent run ID for a chain run. If there is no parent, the run is the root.
+        """
+        if parent_run_id is not None:
+            self._parent_tree[run_id] = parent_run_id
+
+    def _pop_parent_of_run(self, run_id: UUID):
+        """
+        Remove the parent run ID for a chain run.
+        """
+        try:
+            self._parent_tree.pop(run_id)
+        except KeyError:
+            pass
+
+    def _find_root_run(self, run_id: UUID) -> UUID:
+        """
+        Finds the root ID of a chain run.
+        """
+        id: UUID = run_id
+        while id in self._parent_tree:
+            id = self._parent_tree[id]
+        return id
+
+    def _set_run_metadata(
+        self,
+        serialized: Dict[str, Any],
+        run_id: UUID,
+        messages: Union[List[Dict[str, Any]], List[str]],
+        metadata: Optional[Dict[str, Any]] = None,
+        invocation_params: Optional[Dict[str, Any]] = None,
+        **kwargs,
+    ):
+        run: RunMetadata = {
+            "messages": messages,
+            "start_time": time.time(),
+        }
+        if isinstance(invocation_params, dict):
+            run["model_params"] = get_model_params(invocation_params)
+        if isinstance(metadata, dict):
+            if model := metadata.get("ls_model_name"):
+                run["model"] = model
+            if provider := metadata.get("ls_provider"):
+                run["provider"] = provider
+        try:
+            base_url = serialized["kwargs"]["openai_api_base"]
+            if base_url is not None:
+                run["base_url"] = base_url
+        except KeyError:
+            pass
+        self._runs[run_id] = run
+
+    def _pop_run_metadata(self, run_id: UUID) -> Optional[RunMetadata]:
+        end_time = time.time()
+        try:
+            run = self._runs.pop(run_id)
+        except KeyError:
+            log.warning(f"No run metadata found for run {run_id}")
+            return None
+        run["end_time"] = end_time
+        return run
+
+    def _get_trace_id(self, run_id: UUID):
+        trace_id = self._trace_id or self._find_root_run(run_id)
+        if not trace_id:
+            trace_id = uuid.uuid4()
+        return trace_id
+
+
+def _extract_raw_esponse(last_response):
+    """Extract the response from the last response of the LLM call."""
+    # We return the text of the response if not empty
+    if last_response.text is not None and last_response.text.strip() != "":
+        return last_response.text.strip()
+    elif hasattr(last_response, "message"):
+        # Additional kwargs contains the response in case of tool usage
+        return last_response.message.additional_kwargs
+    else:
+        # Not tool usage, some LLM responses can be simply empty
+        return ""
+
+
+def _convert_message_to_dict(message: BaseMessage) -> Dict[str, Any]:
+    # assistant message
+    if isinstance(message, HumanMessage):
+        message_dict = {"role": "user", "content": message.content}
+    elif isinstance(message, AIMessage):
+        message_dict = {"role": "assistant", "content": message.content}
+    elif isinstance(message, SystemMessage):
+        message_dict = {"role": "system", "content": message.content}
+    elif isinstance(message, ToolMessage):
+        message_dict = {"role": "tool", "content": message.content}
+    elif isinstance(message, FunctionMessage):
+        message_dict = {"role": "function", "content": message.content}
+    else:
+        message_dict = {"role": message.type, "content": str(message.content)}
+
+    if "name" in message.additional_kwargs:
+        message_dict["name"] = message.additional_kwargs["name"]
+    if message.additional_kwargs:
+        message_dict["additional_kwargs"] = message.additional_kwargs
+
+    return message_dict
+
+
+def _parse_usage_model(usage: Union[BaseModel, Dict]) -> Tuple[Union[int, None], Union[int, None]]:
+    if isinstance(usage, BaseModel):
+        usage = usage.__dict__
+
+    conversion_list = [
+        # https://pypi.org/project/langchain-anthropic/ (works also for Bedrock-Anthropic)
+        ("input_tokens", "input"),
+        ("output_tokens", "output"),
+        # https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/get-token-count
+        ("prompt_token_count", "input"),
+        ("candidates_token_count", "output"),
+        # Bedrock: https://docs.aws.amazon.com/bedrock/latest/userguide/monitoring-cw.html#runtime-cloudwatch-metrics
+        ("inputTokenCount", "input"),
+        ("outputTokenCount", "output"),
+        # langchain-ibm https://pypi.org/project/langchain-ibm/
+        ("input_token_count", "input"),
+        ("generated_token_count", "output"),
+    ]
+
+    parsed_usage = {}
+    for model_key, type_key in conversion_list:
+        if model_key in usage:
+            captured_count = usage[model_key]
+            final_count = (
+                sum(captured_count) if isinstance(captured_count, list) else captured_count
+            )  # For Bedrock, the token count is a list when streamed
+
+            parsed_usage[type_key] = final_count
+
+    return parsed_usage.get("input"), parsed_usage.get("output")
+
+
+def _parse_usage(response: LLMResult):
+    # langchain-anthropic uses the usage field
+    llm_usage_keys = ["token_usage", "usage"]
+    llm_usage: Tuple[Union[int, None], Union[int, None]] = (None, None)
+    if response.llm_output is not None:
+        for key in llm_usage_keys:
+            if response.llm_output.get(key):
+                llm_usage = _parse_usage_model(response.llm_output[key])
+                break
+
+    if hasattr(response, "generations"):
+        for generation in response.generations:
+            for generation_chunk in generation:
+                if generation_chunk.generation_info and ("usage_metadata" in generation_chunk.generation_info):
+                    llm_usage = _parse_usage_model(generation_chunk.generation_info["usage_metadata"])
+                    break
+
+                message_chunk = getattr(generation_chunk, "message", {})
+                response_metadata = getattr(message_chunk, "response_metadata", {})
+
+                bedrock_anthropic_usage = (
+                    response_metadata.get("usage", None)  # for Bedrock-Anthropic
+                    if isinstance(response_metadata, dict)
+                    else None
+                )
+                bedrock_titan_usage = (
+                    response_metadata.get("amazon-bedrock-invocationMetrics", None)  # for Bedrock-Titan
+                    if isinstance(response_metadata, dict)
+                    else None
+                )
+                ollama_usage = getattr(message_chunk, "usage_metadata", None)  # for Ollama
+
+                chunk_usage = bedrock_anthropic_usage or bedrock_titan_usage or ollama_usage
+                if chunk_usage:
+                    llm_usage = _parse_usage_model(chunk_usage)
+                    break
+
+    return llm_usage
+
+
+def _get_http_status(error: BaseException) -> int:
+    # OpenAI: https://github.com/openai/openai-python/blob/main/src/openai/_exceptions.py
+    # Anthropic: https://github.com/anthropics/anthropic-sdk-python/blob/main/src/anthropic/_exceptions.py
+    # Google: https://github.com/googleapis/python-api-core/blob/main/google/api_core/exceptions.py
+    status_code = getattr(error, "status_code", getattr(error, "code", 0))
+    return status_code
@@ -0,0 +1,4 @@
+from .openai import OpenAI
+from .openai_async import AsyncOpenAI
+
+__all__ = ["OpenAI", "AsyncOpenAI"]
@@ -0,0 +1,237 @@
+import time
+import uuid
+from typing import Any, Dict, Optional
+
+try:
+    import openai
+    import openai.resources
+except ImportError:
+    raise ModuleNotFoundError("Please install the OpenAI SDK to use this feature: 'pip install openai'")
+
+from posthog.ai.utils import call_llm_and_track_usage, get_model_params
+from posthog.client import Client as PostHogClient
+
+
+class OpenAI(openai.OpenAI):
+    """
+    A wrapper around the OpenAI SDK that automatically sends LLM usage events to PostHog.
+    """
+
+    _ph_client: PostHogClient
+
+    def __init__(self, posthog_client: PostHogClient, **kwargs):
+        """
+        Args:
+            api_key: OpenAI API key.
+            posthog_client: If provided, events will be captured via this client instead
+                            of the global posthog.
+            **openai_config: Any additional keyword args to set on openai (e.g. organization="xxx").
+        """
+        super().__init__(**kwargs)
+        self._ph_client = posthog_client
+        self.chat = WrappedChat(self)
+        self.embeddings = WrappedEmbeddings(self)
+
+
+class WrappedChat(openai.resources.chat.Chat):
+    _client: OpenAI
+
+    @property
+    def completions(self):
+        return WrappedCompletions(self._client)
+
+
+class WrappedCompletions(openai.resources.chat.completions.Completions):
+    _client: OpenAI
+
+    def create(
+        self,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        if posthog_trace_id is None:
+            posthog_trace_id = uuid.uuid4()
+
+        if kwargs.get("stream", False):
+            return self._create_streaming(
+                posthog_distinct_id,
+                posthog_trace_id,
+                posthog_properties,
+                **kwargs,
+            )
+
+        return call_llm_and_track_usage(
+            posthog_distinct_id,
+            self._client._ph_client,
+            posthog_trace_id,
+            posthog_properties,
+            self._client.base_url,
+            super().create,
+            **kwargs,
+        )
+
+    def _create_streaming(
+        self,
+        posthog_distinct_id: Optional[str],
+        posthog_trace_id: Optional[str],
+        posthog_properties: Optional[Dict[str, Any]],
+        **kwargs: Any,
+    ):
+        start_time = time.time()
+        usage_stats: Dict[str, int] = {}
+        accumulated_content = []
+        if "stream_options" not in kwargs:
+            kwargs["stream_options"] = {}
+        kwargs["stream_options"]["include_usage"] = True
+        response = super().create(**kwargs)
+
+        def generator():
+            nonlocal usage_stats
+            nonlocal accumulated_content
+            try:
+                for chunk in response:
+                    if hasattr(chunk, "usage") and chunk.usage:
+                        usage_stats = {
+                            k: getattr(chunk.usage, k, 0)
+                            for k in [
+                                "prompt_tokens",
+                                "completion_tokens",
+                                "total_tokens",
+                            ]
+                        }
+
+                    if hasattr(chunk, "choices") and chunk.choices and len(chunk.choices) > 0:
+                        content = chunk.choices[0].delta.content
+                        if content:
+                            accumulated_content.append(content)
+
+                    yield chunk
+
+            finally:
+                end_time = time.time()
+                latency = end_time - start_time
+                output = "".join(accumulated_content)
+                self._capture_streaming_event(
+                    posthog_distinct_id,
+                    posthog_trace_id,
+                    posthog_properties,
+                    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]],
+        kwargs: Dict[str, Any],
+        usage_stats: Dict[str, int],
+        latency: float,
+        output: str,
+    ):
+        if posthog_trace_id is None:
+            posthog_trace_id = uuid.uuid4()
+
+        event_properties = {
+            "$ai_provider": "openai",
+            "$ai_model": kwargs.get("model"),
+            "$ai_model_parameters": get_model_params(kwargs),
+            "$ai_input": kwargs.get("messages"),
+            "$ai_output": {
+                "choices": [
+                    {
+                        "content": output,
+                        "role": "assistant",
+                    }
+                ]
+            },
+            "$ai_http_status": 200,
+            "$ai_input_tokens": usage_stats.get("prompt_tokens", 0),
+            "$ai_output_tokens": usage_stats.get("completion_tokens", 0),
+            "$ai_latency": latency,
+            "$ai_trace_id": posthog_trace_id,
+            "$ai_base_url": str(self._client.base_url),
+            **posthog_properties,
+        }
+
+        if posthog_distinct_id is None:
+            event_properties["$process_person_profile"] = False
+
+        if hasattr(self._client._ph_client, "capture"):
+            self._client._ph_client.capture(
+                distinct_id=posthog_distinct_id or posthog_trace_id,
+                event="$ai_generation",
+                properties=event_properties,
+            )
+
+
+class WrappedEmbeddings(openai.resources.embeddings.Embeddings):
+    _client: OpenAI
+
+    def create(
+        self,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        """
+        Create an embedding using OpenAI's 'embeddings.create' method, but also track usage in PostHog.
+
+        Args:
+            posthog_distinct_id: Optional ID to associate with the usage event.
+            posthog_trace_id: Optional trace UUID for linking events.
+            posthog_properties: Optional dictionary of extra properties to include in the event.
+            **kwargs: Any additional parameters for the OpenAI Embeddings API.
+
+        Returns:
+            The response from OpenAI's embeddings.create call.
+        """
+        if posthog_trace_id is None:
+            posthog_trace_id = uuid.uuid4()
+
+        start_time = time.time()
+        response = super().create(**kwargs)
+        end_time = time.time()
+
+        # Extract usage statistics if available
+        usage_stats = {}
+        if hasattr(response, "usage") and response.usage:
+            usage_stats = {
+                "prompt_tokens": getattr(response.usage, "prompt_tokens", 0),
+                "total_tokens": getattr(response.usage, "total_tokens", 0),
+            }
+
+        latency = end_time - start_time
+
+        # Build the event properties
+        event_properties = {
+            "$ai_provider": "openai",
+            "$ai_model": kwargs.get("model"),
+            "$ai_input": kwargs.get("input"),
+            "$ai_http_status": 200,
+            "$ai_input_tokens": usage_stats.get("prompt_tokens", 0),
+            "$ai_latency": latency,
+            "$ai_trace_id": posthog_trace_id,
+            "$ai_base_url": str(self._client.base_url),
+            **posthog_properties,
+        }
+
+        if posthog_distinct_id is None:
+            event_properties["$process_person_profile"] = False
+
+        # Send capture event for embeddings
+        if hasattr(self._client._ph_client, "capture"):
+            self._client._ph_client.capture(
+                distinct_id=posthog_distinct_id or posthog_trace_id,
+                event="$ai_embedding",
+                properties=event_properties,
+            )
+
+        return response
@@ -0,0 +1,236 @@
+import time
+import uuid
+from typing import Any, Dict, Optional
+
+try:
+    import openai
+    import openai.resources
+except ImportError:
+    raise ModuleNotFoundError("Please install the OpenAI SDK to use this feature: 'pip install openai'")
+
+from posthog.ai.utils import call_llm_and_track_usage_async, get_model_params
+from posthog.client import Client as PostHogClient
+
+
+class AsyncOpenAI(openai.AsyncOpenAI):
+    """
+    An async wrapper around the OpenAI SDK that automatically sends LLM usage events to PostHog.
+    """
+
+    _ph_client: PostHogClient
+
+    def __init__(self, posthog_client: PostHogClient, **kwargs):
+        """
+        Args:
+            api_key: OpenAI API key.
+            posthog_client: If provided, events will be captured via this client instance.
+            **openai_config: Additional keyword args (e.g. organization="xxx").
+        """
+        super().__init__(**kwargs)
+        self._ph_client = posthog_client
+        self.chat = WrappedChat(self)
+        self.embeddings = WrappedEmbeddings(self)
+
+
+class WrappedChat(openai.resources.chat.AsyncChat):
+    _client: AsyncOpenAI
+
+    @property
+    def completions(self):
+        return WrappedCompletions(self._client)
+
+
+class WrappedCompletions(openai.resources.chat.completions.AsyncCompletions):
+    _client: AsyncOpenAI
+
+    async def create(
+        self,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        if posthog_trace_id is None:
+            posthog_trace_id = uuid.uuid4()
+
+        # If streaming, handle streaming specifically
+        if kwargs.get("stream", False):
+            return await self._create_streaming(
+                posthog_distinct_id,
+                posthog_trace_id,
+                posthog_properties,
+                **kwargs,
+            )
+
+        response = await call_llm_and_track_usage_async(
+            posthog_distinct_id,
+            self._client._ph_client,
+            posthog_trace_id,
+            posthog_properties,
+            self._client.base_url,
+            super().create,
+            **kwargs,
+        )
+        return response
+
+    async def _create_streaming(
+        self,
+        posthog_distinct_id: Optional[str],
+        posthog_trace_id: Optional[str],
+        posthog_properties: Optional[Dict[str, Any]],
+        **kwargs: Any,
+    ):
+        start_time = time.time()
+        usage_stats: Dict[str, int] = {}
+        accumulated_content = []
+        if "stream_options" not in kwargs:
+            kwargs["stream_options"] = {}
+        kwargs["stream_options"]["include_usage"] = True
+        response = await super().create(**kwargs)
+
+        async def async_generator():
+            nonlocal usage_stats, accumulated_content
+            try:
+                async for chunk in response:
+                    if hasattr(chunk, "usage") and chunk.usage:
+                        usage_stats = {
+                            k: getattr(chunk.usage, k, 0)
+                            for k in [
+                                "prompt_tokens",
+                                "completion_tokens",
+                                "total_tokens",
+                            ]
+                        }
+                    if hasattr(chunk, "choices") and chunk.choices and len(chunk.choices) > 0:
+                        content = chunk.choices[0].delta.content
+                        if content:
+                            accumulated_content.append(content)
+
+                    yield chunk
+
+            finally:
+                end_time = time.time()
+                latency = end_time - start_time
+                output = "".join(accumulated_content)
+                self._capture_streaming_event(
+                    posthog_distinct_id,
+                    posthog_trace_id,
+                    posthog_properties,
+                    kwargs,
+                    usage_stats,
+                    latency,
+                    output,
+                )
+
+        return async_generator()
+
+    def _capture_streaming_event(
+        self,
+        posthog_distinct_id: Optional[str],
+        posthog_trace_id: Optional[str],
+        posthog_properties: Optional[Dict[str, Any]],
+        kwargs: Dict[str, Any],
+        usage_stats: Dict[str, int],
+        latency: float,
+        output: str,
+    ):
+        if posthog_trace_id is None:
+            posthog_trace_id = uuid.uuid4()
+
+        event_properties = {
+            "$ai_provider": "openai",
+            "$ai_model": kwargs.get("model"),
+            "$ai_model_parameters": get_model_params(kwargs),
+            "$ai_input": kwargs.get("messages"),
+            "$ai_output": {
+                "choices": [
+                    {
+                        "content": output,
+                        "role": "assistant",
+                    }
+                ]
+            },
+            "$ai_http_status": 200,
+            "$ai_input_tokens": usage_stats.get("prompt_tokens", 0),
+            "$ai_output_tokens": usage_stats.get("completion_tokens", 0),
+            "$ai_latency": latency,
+            "$ai_trace_id": posthog_trace_id,
+            "$ai_base_url": str(self._client.base_url),
+            **posthog_properties,
+        }
+
+        if posthog_distinct_id is None:
+            event_properties["$process_person_profile"] = False
+
+        if hasattr(self._client._ph_client, "capture"):
+            self._client._ph_client.capture(
+                distinct_id=posthog_distinct_id or posthog_trace_id,
+                event="$ai_generation",
+                properties=event_properties,
+            )
+
+
+class WrappedEmbeddings(openai.resources.embeddings.AsyncEmbeddings):
+    _client: AsyncOpenAI
+
+    async def create(
+        self,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        """
+        Create an embedding using OpenAI's 'embeddings.create' method, but also track usage in PostHog.
+
+        Args:
+            posthog_distinct_id: Optional ID to associate with the usage event.
+            posthog_trace_id: Optional trace UUID for linking events.
+            posthog_properties: Optional dictionary of extra properties to include in the event.
+            **kwargs: Any additional parameters for the OpenAI Embeddings API.
+
+        Returns:
+            The response from OpenAI's embeddings.create call.
+        """
+        if posthog_trace_id is None:
+            posthog_trace_id = uuid.uuid4()
+
+        start_time = time.time()
+        response = await super().create(**kwargs)
+        end_time = time.time()
+
+        # Extract usage statistics if available
+        usage_stats = {}
+        if hasattr(response, "usage") and response.usage:
+            usage_stats = {
+                "prompt_tokens": getattr(response.usage, "prompt_tokens", 0),
+                "total_tokens": getattr(response.usage, "total_tokens", 0),
+            }
+
+        latency = end_time - start_time
+
+        # Build the event properties
+        event_properties = {
+            "$ai_provider": "openai",
+            "$ai_model": kwargs.get("model"),
+            "$ai_input": kwargs.get("input"),
+            "$ai_http_status": 200,
+            "$ai_input_tokens": usage_stats.get("prompt_tokens", 0),
+            "$ai_latency": latency,
+            "$ai_trace_id": posthog_trace_id,
+            "$ai_base_url": str(self._client.base_url),
+            **posthog_properties,
+        }
+
+        if posthog_distinct_id is None:
+            event_properties["$process_person_profile"] = False
+
+        # Send capture event for embeddings
+        if hasattr(self._client._ph_client, "capture"):
+            self._client._ph_client.capture(
+                distinct_id=posthog_distinct_id or posthog_trace_id,
+                event="$ai_embedding",
+                properties=event_properties,
+            )
+
+        return response
@@ -0,0 +1,178 @@
+import time
+import uuid
+from typing import Any, Callable, Dict, Optional
+
+from httpx import URL
+
+from posthog.client import Client as PostHogClient
+
+
+def get_model_params(kwargs: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Extracts model parameters from the kwargs dictionary.
+    """
+    model_params = {}
+    for param in [
+        "temperature",
+        "max_tokens",  # Deprecated field
+        "max_completion_tokens",
+        "top_p",
+        "frequency_penalty",
+        "presence_penalty",
+        "n",
+        "stop",
+        "stream",
+    ]:
+        if param in kwargs and kwargs[param] is not None:
+            model_params[param] = kwargs[param]
+    return model_params
+
+
+def format_response(response):
+    """
+    Format a regular (non-streaming) response.
+    """
+    output = {"choices": []}
+    if response is None:
+        return output
+    for choice in response.choices:
+        if choice.message.content:
+            output["choices"].append(
+                {
+                    "content": choice.message.content,
+                    "role": choice.message.role,
+                }
+            )
+    return output
+
+
+def call_llm_and_track_usage(
+    posthog_distinct_id: Optional[str],
+    ph_client: PostHogClient,
+    posthog_trace_id: Optional[str],
+    posthog_properties: Optional[Dict[str, Any]],
+    base_url: URL,
+    call_method: Callable[..., Any],
+    **kwargs: Any,
+) -> Any:
+    """
+    Common usage-tracking logic for both sync and async calls.
+    call_method: the llm call method (e.g. openai.chat.completions.create)
+    """
+    start_time = time.time()
+    response = None
+    error = None
+    http_status = 200
+    usage: Dict[str, Any] = {}
+
+    try:
+        response = call_method(**kwargs)
+    except Exception as exc:
+        error = exc
+        http_status = getattr(exc, "status_code", 0)  # default to 0 becuase its likely an SDK error
+    finally:
+        end_time = time.time()
+        latency = end_time - start_time
+
+        if posthog_trace_id is None:
+            posthog_trace_id = uuid.uuid4()
+
+        if response and hasattr(response, "usage"):
+            usage = response.usage.model_dump()
+
+        input_tokens = usage.get("prompt_tokens", 0)
+        output_tokens = usage.get("completion_tokens", 0)
+        event_properties = {
+            "$ai_provider": "openai",
+            "$ai_model": kwargs.get("model"),
+            "$ai_model_parameters": get_model_params(kwargs),
+            "$ai_input": kwargs.get("messages"),
+            "$ai_output": format_response(response),
+            "$ai_http_status": http_status,
+            "$ai_input_tokens": input_tokens,
+            "$ai_output_tokens": output_tokens,
+            "$ai_latency": latency,
+            "$ai_trace_id": posthog_trace_id,
+            "$ai_base_url": str(base_url),
+            **(posthog_properties or {}),
+        }
+
+        if posthog_distinct_id is None:
+            event_properties["$process_person_profile"] = False
+
+        # send the event to posthog
+        if hasattr(ph_client, "capture") and callable(ph_client.capture):
+            ph_client.capture(
+                distinct_id=posthog_distinct_id or posthog_trace_id,
+                event="$ai_generation",
+                properties=event_properties,
+            )
+
+    if error:
+        raise error
+
+    return response
+
+
+async def call_llm_and_track_usage_async(
+    posthog_distinct_id: Optional[str],
+    ph_client: PostHogClient,
+    posthog_trace_id: Optional[str],
+    posthog_properties: Optional[Dict[str, Any]],
+    base_url: URL,
+    call_async_method: Callable[..., Any],
+    **kwargs: Any,
+) -> Any:
+    start_time = time.time()
+    response = None
+    error = None
+    http_status = 200
+    usage: Dict[str, Any] = {}
+
+    try:
+        response = await call_async_method(**kwargs)
+    except Exception as exc:
+        error = exc
+        http_status = getattr(exc, "status_code", 0)  # default to 0 because its likely an SDK error
+    finally:
+        end_time = time.time()
+        latency = end_time - start_time
+
+        if posthog_trace_id is None:
+            posthog_trace_id = uuid.uuid4()
+
+        if response and hasattr(response, "usage"):
+            usage = response.usage.model_dump()
+
+        input_tokens = usage.get("prompt_tokens", 0)
+        output_tokens = usage.get("completion_tokens", 0)
+        event_properties = {
+            "$ai_provider": "openai",
+            "$ai_model": kwargs.get("model"),
+            "$ai_model_parameters": get_model_params(kwargs),
+            "$ai_input": kwargs.get("messages"),
+            "$ai_output": format_response(response),
+            "$ai_http_status": http_status,
+            "$ai_input_tokens": input_tokens,
+            "$ai_output_tokens": output_tokens,
+            "$ai_latency": latency,
+            "$ai_trace_id": posthog_trace_id,
+            "$ai_base_url": str(base_url),
+            **(posthog_properties or {}),
+        }
+
+        if posthog_distinct_id is None:
+            event_properties["$process_person_profile"] = False
+
+        # send the event to posthog
+        if hasattr(ph_client, "capture") and callable(ph_client.capture):
+            ph_client.capture(
+                distinct_id=posthog_distinct_id or posthog_trace_id,
+                event="$ai_generation",
+                properties=event_properties,
+            )
+
+    if error:
+        raise error
+
+    return response
@@ -304,6 +304,7 @@ class Client(object):
        timestamp=None,
        uuid=None,
        disable_geoip=None,
+        distinct_id=None,
    ):
        properties = properties or {}
        context = context or {}
@@ -311,6 +312,11 @@ class Client(object):
        require("group_key", group_key, ID_TYPES)
        require("properties", properties, dict)

+        if distinct_id:
+            require("distinct_id", distinct_id, ID_TYPES)
+        else:
+            distinct_id = "${}_{}".format(group_type, group_key)
+
        msg = {
            "event": "$groupidentify",
            "properties": {
@@ -318,7 +324,7 @@ class Client(object):
                "$group_key": group_key,
                "$group_set": properties,
            },
-            "distinct_id": "${}_{}".format(group_type, group_key),
+            "distinct_id": distinct_id,
            "timestamp": timestamp,
            "context": context,
            "uuid": uuid,
@@ -0,0 +1,4 @@
+import pytest
+
+pytest.importorskip("langchain")
+pytest.importorskip("langchain_community")
@@ -0,0 +1,597 @@
+import math
+import os
+import time
+import uuid
+from unittest.mock import patch
+
+import pytest
+from langchain_community.chat_models.fake import FakeMessagesListChatModel
+from langchain_community.llms.fake import FakeListLLM, FakeStreamingListLLM
+from langchain_core.messages import AIMessage
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.runnables import RunnableLambda
+from langchain_openai.chat_models import ChatOpenAI
+
+from posthog.ai.langchain import CallbackHandler
+
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+
+
+@pytest.fixture(scope="function")
+def mock_client():
+    with patch("posthog.client.Client") as mock_client:
+        yield mock_client
+
+
+def test_parent_capture(mock_client):
+    callbacks = CallbackHandler(mock_client)
+    parent_run_id = uuid.uuid4()
+    run_id = uuid.uuid4()
+    callbacks._set_parent_of_run(run_id, parent_run_id)
+    assert callbacks._parent_tree == {run_id: parent_run_id}
+    callbacks._pop_parent_of_run(run_id)
+    assert callbacks._parent_tree == {}
+    callbacks._pop_parent_of_run(parent_run_id)  # should not raise
+
+
+def test_find_root_run(mock_client):
+    callbacks = CallbackHandler(mock_client)
+    root_run_id = uuid.uuid4()
+    parent_run_id = uuid.uuid4()
+    run_id = uuid.uuid4()
+    callbacks._set_parent_of_run(run_id, parent_run_id)
+    callbacks._set_parent_of_run(parent_run_id, root_run_id)
+    assert callbacks._find_root_run(run_id) == root_run_id
+    new_run_id = uuid.uuid4()
+    assert callbacks._find_root_run(new_run_id) == new_run_id
+
+
+def test_trace_id_generation(mock_client):
+    callbacks = CallbackHandler(mock_client)
+    run_id = uuid.uuid4()
+    with patch("uuid.uuid4", return_value=run_id):
+        assert callbacks._get_trace_id(run_id) == run_id
+    run_id = uuid.uuid4()
+    callbacks = CallbackHandler(mock_client, trace_id=run_id)
+    assert callbacks._get_trace_id(uuid.uuid4()) == run_id
+
+
+def test_metadata_capture(mock_client):
+    callbacks = CallbackHandler(mock_client)
+    run_id = uuid.uuid4()
+    with patch("time.time", return_value=1234567890):
+        callbacks._set_run_metadata(
+            {"kwargs": {"openai_api_base": "https://us.posthog.com"}},
+            run_id,
+            messages=[{"role": "user", "content": "Who won the world series in 2020?"}],
+            invocation_params={"temperature": 0.5},
+            metadata={"ls_model_name": "hog-mini", "ls_provider": "posthog"},
+        )
+    expected = {
+        "model": "hog-mini",
+        "messages": [{"role": "user", "content": "Who won the world series in 2020?"}],
+        "start_time": 1234567890,
+        "model_params": {"temperature": 0.5},
+        "provider": "posthog",
+        "base_url": "https://us.posthog.com",
+    }
+    assert callbacks._runs[run_id] == expected
+    with patch("time.time", return_value=1234567891):
+        run = callbacks._pop_run_metadata(run_id)
+    assert run == {**expected, "end_time": 1234567891}
+    assert callbacks._runs == {}
+    callbacks._pop_run_metadata(uuid.uuid4())  # should not raise
+
+
+@pytest.mark.parametrize("stream", [True, False])
+def test_basic_chat_chain(mock_client, stream):
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", "You are a helpful assistant."),
+            ("user", "Who won the world series in 2020?"),
+        ]
+    )
+    model = FakeMessagesListChatModel(
+        responses=[
+            AIMessage(
+                content="The Los Angeles Dodgers won the World Series in 2020.",
+                usage_metadata={"input_tokens": 10, "output_tokens": 10, "total_tokens": 20},
+            )
+        ]
+    )
+    callbacks = [CallbackHandler(mock_client)]
+    chain = prompt | model
+    if stream:
+        result = [m for m in chain.stream({}, config={"callbacks": callbacks})][0]
+    else:
+        result = chain.invoke({}, config={"callbacks": callbacks})
+
+    assert result.content == "The Los Angeles Dodgers won the World Series in 2020."
+    assert mock_client.capture.call_count == 1
+    args = mock_client.capture.call_args[1]
+    props = args["properties"]
+
+    assert args["event"] == "$ai_generation"
+    assert "distinct_id" in args
+    assert "$ai_model" in props
+    assert "$ai_provider" in props
+    assert props["$ai_input"] == [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": "Who won the world series in 2020?"},
+    ]
+    assert props["$ai_output"] == {
+        "choices": [{"role": "assistant", "content": "The Los Angeles Dodgers won the World Series in 2020."}]
+    }
+    assert props["$ai_input_tokens"] == 10
+    assert props["$ai_output_tokens"] == 10
+    assert props["$ai_http_status"] == 200
+    assert props["$ai_trace_id"] is not None
+    assert isinstance(props["$ai_latency"], float)
+
+
+@pytest.mark.parametrize("stream", [True, False])
+async def test_async_basic_chat_chain(mock_client, stream):
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", "You are a helpful assistant."),
+            ("user", "Who won the world series in 2020?"),
+        ]
+    )
+    model = FakeMessagesListChatModel(
+        responses=[
+            AIMessage(
+                content="The Los Angeles Dodgers won the World Series in 2020.",
+                usage_metadata={"input_tokens": 10, "output_tokens": 10, "total_tokens": 20},
+            )
+        ]
+    )
+    callbacks = [CallbackHandler(mock_client)]
+    chain = prompt | model
+    if stream:
+        result = [m async for m in chain.astream({}, config={"callbacks": callbacks})][0]
+    else:
+        result = await chain.ainvoke({}, config={"callbacks": callbacks})
+    assert result.content == "The Los Angeles Dodgers won the World Series in 2020."
+    assert mock_client.capture.call_count == 1
+
+    args = mock_client.capture.call_args[1]
+    props = args["properties"]
+    assert args["event"] == "$ai_generation"
+    assert "distinct_id" in args
+    assert "$ai_model" in props
+    assert "$ai_provider" in props
+    assert props["$ai_input"] == [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": "Who won the world series in 2020?"},
+    ]
+    assert props["$ai_output"] == {
+        "choices": [{"role": "assistant", "content": "The Los Angeles Dodgers won the World Series in 2020."}]
+    }
+    assert props["$ai_input_tokens"] == 10
+    assert props["$ai_output_tokens"] == 10
+    assert props["$ai_http_status"] == 200
+    assert props["$ai_trace_id"] is not None
+    assert isinstance(props["$ai_latency"], float)
+
+
+@pytest.mark.parametrize(
+    "Model,stream",
+    [(FakeListLLM, True), (FakeListLLM, False), (FakeStreamingListLLM, True), (FakeStreamingListLLM, False)],
+)
+def test_basic_llm_chain(mock_client, Model, stream):
+    model = Model(responses=["The Los Angeles Dodgers won the World Series in 2020."])
+    callbacks: list[CallbackHandler] = [CallbackHandler(mock_client)]
+
+    if stream:
+        result = "".join(
+            [m for m in model.stream("Who won the world series in 2020?", config={"callbacks": callbacks})]
+        )
+    else:
+        result = model.invoke("Who won the world series in 2020?", config={"callbacks": callbacks})
+    assert result == "The Los Angeles Dodgers won the World Series in 2020."
+
+    assert mock_client.capture.call_count == 1
+    args = mock_client.capture.call_args[1]
+    props = args["properties"]
+
+    assert args["event"] == "$ai_generation"
+    assert "distinct_id" in args
+    assert "$ai_model" in props
+    assert "$ai_provider" in props
+    assert props["$ai_input"] == ["Who won the world series in 2020?"]
+    assert props["$ai_output"] == {"choices": ["The Los Angeles Dodgers won the World Series in 2020."]}
+    assert props["$ai_http_status"] == 200
+    assert props["$ai_trace_id"] is not None
+    assert isinstance(props["$ai_latency"], float)
+
+
+@pytest.mark.parametrize(
+    "Model,stream",
+    [(FakeListLLM, True), (FakeListLLM, False), (FakeStreamingListLLM, True), (FakeStreamingListLLM, False)],
+)
+async def test_async_basic_llm_chain(mock_client, Model, stream):
+    model = Model(responses=["The Los Angeles Dodgers won the World Series in 2020."])
+    callbacks: list[CallbackHandler] = [CallbackHandler(mock_client)]
+
+    if stream:
+        result = "".join(
+            [m async for m in model.astream("Who won the world series in 2020?", config={"callbacks": callbacks})]
+        )
+    else:
+        result = await model.ainvoke("Who won the world series in 2020?", config={"callbacks": callbacks})
+    assert result == "The Los Angeles Dodgers won the World Series in 2020."
+
+    assert mock_client.capture.call_count == 1
+    args = mock_client.capture.call_args[1]
+    props = args["properties"]
+
+    assert args["event"] == "$ai_generation"
+    assert "distinct_id" in args
+    assert "$ai_model" in props
+    assert "$ai_provider" in props
+    assert props["$ai_input"] == ["Who won the world series in 2020?"]
+    assert props["$ai_output"] == {"choices": ["The Los Angeles Dodgers won the World Series in 2020."]}
+    assert props["$ai_http_status"] == 200
+    assert props["$ai_trace_id"] is not None
+    assert isinstance(props["$ai_latency"], float)
+
+
+def test_trace_id_for_multiple_chains(mock_client):
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("user", "Foo"),
+        ]
+    )
+    model = FakeMessagesListChatModel(responses=[AIMessage(content="Bar")])
+    callbacks = [CallbackHandler(mock_client)]
+    chain = prompt | model | RunnableLambda(lambda x: [x]) | model
+    result = chain.invoke({}, config={"callbacks": callbacks})
+
+    assert result.content == "Bar"
+    assert mock_client.capture.call_count == 2
+
+    first_call_args = mock_client.capture.call_args_list[0][1]
+    first_call_props = first_call_args["properties"]
+    assert first_call_args["event"] == "$ai_generation"
+    assert "distinct_id" in first_call_args
+    assert "$ai_model" in first_call_props
+    assert "$ai_provider" in first_call_props
+    assert first_call_props["$ai_input"] == [{"role": "user", "content": "Foo"}]
+    assert first_call_props["$ai_output"] == {"choices": [{"role": "assistant", "content": "Bar"}]}
+    assert first_call_props["$ai_http_status"] == 200
+    assert first_call_props["$ai_trace_id"] is not None
+    assert isinstance(first_call_props["$ai_latency"], float)
+
+    second_call_args = mock_client.capture.call_args_list[1][1]
+    second_call_props = second_call_args["properties"]
+    assert second_call_args["event"] == "$ai_generation"
+    assert "distinct_id" in second_call_args
+    assert "$ai_model" in second_call_props
+    assert "$ai_provider" in second_call_props
+    assert second_call_props["$ai_input"] == [{"role": "assistant", "content": "Bar"}]
+    assert second_call_props["$ai_output"] == {"choices": [{"role": "assistant", "content": "Bar"}]}
+    assert second_call_props["$ai_http_status"] == 200
+    assert second_call_props["$ai_trace_id"] is not None
+    assert isinstance(second_call_props["$ai_latency"], float)
+
+    # Check that the trace_id is the same as the first call
+    assert first_call_props["$ai_trace_id"] == second_call_props["$ai_trace_id"]
+
+
+def test_personless_mode(mock_client):
+    prompt = ChatPromptTemplate.from_messages([("user", "Foo")])
+    chain = prompt | FakeMessagesListChatModel(responses=[AIMessage(content="Bar")])
+    chain.invoke({}, config={"callbacks": [CallbackHandler(mock_client)]})
+    assert mock_client.capture.call_count == 1
+    args = mock_client.capture.call_args_list[0][1]
+    assert args["properties"]["$process_person_profile"] is False
+
+    id = uuid.uuid4()
+    chain.invoke({}, config={"callbacks": [CallbackHandler(mock_client, distinct_id=id)]})
+    assert mock_client.capture.call_count == 2
+    args = mock_client.capture.call_args_list[1][1]
+    assert "$process_person_profile" not in args["properties"]
+    assert args["distinct_id"] == id
+
+
+def test_personless_mode_exception(mock_client):
+    prompt = ChatPromptTemplate.from_messages([("user", "Foo")])
+    chain = prompt | ChatOpenAI(api_key="test", model="gpt-4o-mini")
+    callbacks = CallbackHandler(mock_client)
+    with pytest.raises(Exception):
+        chain.invoke({}, config={"callbacks": [callbacks]})
+    assert mock_client.capture.call_count == 1
+    args = mock_client.capture.call_args_list[0][1]
+    assert args["properties"]["$process_person_profile"] is False
+
+    id = uuid.uuid4()
+    with pytest.raises(Exception):
+        chain.invoke({}, config={"callbacks": [CallbackHandler(mock_client, distinct_id=id)]})
+    assert mock_client.capture.call_count == 2
+    args = mock_client.capture.call_args_list[1][1]
+    assert "$process_person_profile" not in args["properties"]
+    assert args["distinct_id"] == id
+
+
+def test_metadata(mock_client):
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("user", "Foo"),
+        ]
+    )
+    model = FakeMessagesListChatModel(responses=[AIMessage(content="Bar")])
+    callbacks = [
+        CallbackHandler(mock_client, trace_id="test-trace-id", distinct_id="test_id", properties={"foo": "bar"})
+    ]
+    chain = prompt | model
+    result = chain.invoke({}, config={"callbacks": callbacks})
+
+    assert result.content == "Bar"
+    assert mock_client.capture.call_count == 1
+
+    first_call_args = mock_client.capture.call_args[1]
+    assert first_call_args["distinct_id"] == "test_id"
+
+    first_call_props = first_call_args["properties"]
+    assert first_call_args["event"] == "$ai_generation"
+    assert first_call_props["$ai_trace_id"] == "test-trace-id"
+    assert first_call_props["foo"] == "bar"
+    assert first_call_props["$ai_input"] == [{"role": "user", "content": "Foo"}]
+    assert first_call_props["$ai_output"] == {"choices": [{"role": "assistant", "content": "Bar"}]}
+    assert first_call_props["$ai_http_status"] == 200
+    assert isinstance(first_call_props["$ai_latency"], float)
+
+
+def test_callbacks_logic(mock_client):
+    prompt = ChatPromptTemplate.from_messages([("user", "Foo")])
+    model = FakeMessagesListChatModel(responses=[AIMessage(content="Bar")])
+    callbacks = CallbackHandler(mock_client, trace_id="test-trace-id", distinct_id="test_id", properties={"foo": "bar"})
+    chain = prompt | model
+
+    chain.invoke({}, config={"callbacks": [callbacks]})
+    assert callbacks._runs == {}
+    assert callbacks._parent_tree == {}
+
+    def assert_intermediary_run(m):
+        assert callbacks._runs == {}
+        assert len(callbacks._parent_tree.items()) == 1
+        return [m]
+
+    (chain | RunnableLambda(assert_intermediary_run) | model).invoke({}, config={"callbacks": [callbacks]})
+    assert callbacks._runs == {}
+    assert callbacks._parent_tree == {}
+
+
+def test_exception_in_chain(mock_client):
+    def runnable(_):
+        raise ValueError("test")
+
+    callbacks = CallbackHandler(mock_client)
+    with pytest.raises(ValueError):
+        RunnableLambda(runnable).invoke({}, config={"callbacks": [callbacks]})
+
+    assert callbacks._runs == {}
+    assert callbacks._parent_tree == {}
+    assert mock_client.capture.call_count == 0
+
+
+def test_openai_error(mock_client):
+    prompt = ChatPromptTemplate.from_messages([("user", "Foo")])
+    chain = prompt | ChatOpenAI(api_key="test", model="gpt-4o-mini")
+    callbacks = CallbackHandler(mock_client)
+
+    # 401
+    with pytest.raises(Exception):
+        chain.invoke({}, config={"callbacks": [callbacks]})
+
+    assert callbacks._runs == {}
+    assert callbacks._parent_tree == {}
+    assert mock_client.capture.call_count == 1
+    args = mock_client.capture.call_args[1]
+    props = args["properties"]
+    assert props["$ai_http_status"] == 401
+    assert props["$ai_input"] == [{"role": "user", "content": "Foo"}]
+    assert "$ai_output" not in props
+
+
+@pytest.mark.skipif(not OPENAI_API_KEY, reason="OpenAI API key not set")
+def test_openai_chain(mock_client):
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", 'You must always answer with "Bar".'),
+            ("user", "Foo"),
+        ]
+    )
+    chain = prompt | ChatOpenAI(
+        api_key=OPENAI_API_KEY,
+        model="gpt-4o-mini",
+        temperature=0,
+        max_tokens=1,
+    )
+    callbacks = CallbackHandler(mock_client, trace_id="test-trace-id", distinct_id="test_id", properties={"foo": "bar"})
+    start_time = time.time()
+    result = chain.invoke({}, config={"callbacks": [callbacks]})
+    approximate_latency = math.floor(time.time() - start_time)
+
+    assert result.content == "Bar"
+    assert mock_client.capture.call_count == 1
+
+    first_call_args = mock_client.capture.call_args[1]
+    first_call_props = first_call_args["properties"]
+    assert first_call_args["event"] == "$ai_generation"
+    assert first_call_props["$ai_trace_id"] == "test-trace-id"
+    assert first_call_props["$ai_provider"] == "openai"
+    assert first_call_props["$ai_model"] == "gpt-4o-mini"
+    assert first_call_props["foo"] == "bar"
+
+    # langchain-openai for langchain v3
+    if "max_completion_tokens" in first_call_props["$ai_model_parameters"]:
+        assert first_call_props["$ai_model_parameters"] == {
+            "temperature": 0.0,
+            "max_completion_tokens": 1,
+            "stream": False,
+        }
+    else:
+        assert first_call_props["$ai_model_parameters"] == {
+            "temperature": 0.0,
+            "max_tokens": 1,
+            "n": 1,
+            "stream": False,
+        }
+    assert first_call_props["$ai_input"] == [
+        {"role": "system", "content": 'You must always answer with "Bar".'},
+        {"role": "user", "content": "Foo"},
+    ]
+    assert first_call_props["$ai_output"] == {
+        "choices": [
+            {
+                "role": "assistant",
+                "content": "Bar",
+                "additional_kwargs": {"refusal": None},
+            }
+        ]
+    }
+    assert first_call_props["$ai_http_status"] == 200
+    assert isinstance(first_call_props["$ai_latency"], float)
+    assert min(approximate_latency - 1, 0) <= math.floor(first_call_props["$ai_latency"]) <= approximate_latency
+    assert first_call_props["$ai_input_tokens"] == 20
+    assert first_call_props["$ai_output_tokens"] == 1
+
+
+@pytest.mark.skipif(not OPENAI_API_KEY, reason="OpenAI API key not set")
+def test_openai_captures_multiple_generations(mock_client):
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", 'You must always answer with "Bar".'),
+            ("user", "Foo"),
+        ]
+    )
+    chain = prompt | ChatOpenAI(
+        api_key=OPENAI_API_KEY,
+        model="gpt-4o-mini",
+        temperature=0,
+        max_tokens=1,
+        n=2,
+    )
+    callbacks = CallbackHandler(mock_client)
+    result = chain.invoke({}, config={"callbacks": [callbacks]})
+
+    assert result.content == "Bar"
+    assert mock_client.capture.call_count == 1
+
+    first_call_args = mock_client.capture.call_args[1]
+    first_call_props = first_call_args["properties"]
+    assert first_call_props["$ai_input"] == [
+        {"role": "system", "content": 'You must always answer with "Bar".'},
+        {"role": "user", "content": "Foo"},
+    ]
+    assert first_call_props["$ai_output"] == {
+        "choices": [
+            {
+                "role": "assistant",
+                "content": "Bar",
+                "additional_kwargs": {"refusal": None},
+            },
+            {
+                "role": "assistant",
+                "content": "Bar",
+            },
+        ]
+    }
+
+    # langchain-openai for langchain v3
+    if "max_completion_tokens" in first_call_props["$ai_model_parameters"]:
+        assert first_call_props["$ai_model_parameters"] == {
+            "temperature": 0.0,
+            "max_completion_tokens": 1,
+            "stream": False,
+            "n": 2,
+        }
+    else:
+        assert first_call_props["$ai_model_parameters"] == {
+            "temperature": 0.0,
+            "max_tokens": 1,
+            "stream": False,
+            "n": 2,
+        }
+    assert first_call_props["$ai_http_status"] == 200
+
+
+@pytest.mark.skipif(not OPENAI_API_KEY, reason="OpenAI API key not set")
+def test_openai_streaming(mock_client):
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", 'You must always answer with "Bar".'),
+            ("user", "Foo"),
+        ]
+    )
+    chain = prompt | ChatOpenAI(
+        api_key=OPENAI_API_KEY, model="gpt-4o-mini", temperature=0, max_tokens=1, stream=True, stream_usage=True
+    )
+    callbacks = CallbackHandler(mock_client)
+    result = [m for m in chain.stream({}, config={"callbacks": [callbacks]})]
+    result = sum(result[1:], result[0])
+
+    assert result.content == "Bar"
+    assert mock_client.capture.call_count == 1
+
+    first_call_args = mock_client.capture.call_args[1]
+    first_call_props = first_call_args["properties"]
+
+    assert first_call_props["$ai_model_parameters"]["stream"]
+    assert first_call_props["$ai_input"] == [
+        {"role": "system", "content": 'You must always answer with "Bar".'},
+        {"role": "user", "content": "Foo"},
+    ]
+    assert first_call_props["$ai_output"] == {"choices": [{"role": "assistant", "content": "Bar"}]}
+    assert first_call_props["$ai_http_status"] == 200
+    assert first_call_props["$ai_input_tokens"] == 20
+    assert first_call_props["$ai_output_tokens"] == 1
+
+
+@pytest.mark.skipif(not OPENAI_API_KEY, reason="OpenAI API key not set")
+async def test_async_openai_streaming(mock_client):
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", 'You must always answer with "Bar".'),
+            ("user", "Foo"),
+        ]
+    )
+    chain = prompt | ChatOpenAI(
+        api_key=OPENAI_API_KEY, model="gpt-4o-mini", temperature=0, max_tokens=1, stream=True, stream_usage=True
+    )
+    callbacks = CallbackHandler(mock_client)
+    result = [m async for m in chain.astream({}, config={"callbacks": [callbacks]})]
+    result = sum(result[1:], result[0])
+
+    assert result.content == "Bar"
+    assert mock_client.capture.call_count == 1
+
+    first_call_args = mock_client.capture.call_args[1]
+    first_call_props = first_call_args["properties"]
+
+    assert first_call_props["$ai_model_parameters"]["stream"]
+    assert first_call_props["$ai_input"] == [
+        {"role": "system", "content": 'You must always answer with "Bar".'},
+        {"role": "user", "content": "Foo"},
+    ]
+    assert first_call_props["$ai_output"] == {"choices": [{"role": "assistant", "content": "Bar"}]}
+    assert first_call_props["$ai_http_status"] == 200
+    assert first_call_props["$ai_input_tokens"] == 20
+    assert first_call_props["$ai_output_tokens"] == 1
+
+
+def test_base_url_retrieval(mock_client):
+    prompt = ChatPromptTemplate.from_messages([("user", "Foo")])
+    chain = prompt | ChatOpenAI(
+        api_key="test",
+        model="posthog-mini",
+        base_url="https://test.posthog.com",
+    )
+    callbacks = CallbackHandler(mock_client)
+    with pytest.raises(Exception):
+        chain.invoke({}, config={"callbacks": [callbacks]})
+
+    assert mock_client.capture.call_count == 1
+    call = mock_client.capture.call_args[1]
+    assert call["properties"]["$ai_base_url"] == "https://test.posthog.com"
@@ -0,0 +1,117 @@
+import time
+from unittest.mock import patch
+
+import pytest
+from openai.types.chat import ChatCompletion, ChatCompletionMessage
+from openai.types.chat.chat_completion import Choice
+from openai.types.completion_usage import CompletionUsage
+from openai.types.create_embedding_response import CreateEmbeddingResponse, Usage
+from openai.types.embedding import Embedding
+
+from posthog.ai.openai import OpenAI
+
+
+@pytest.fixture
+def mock_client():
+    with patch("posthog.client.Client") as mock_client:
+        yield mock_client
+
+
+@pytest.fixture
+def mock_openai_response():
+    return ChatCompletion(
+        id="test",
+        model="gpt-4",
+        object="chat.completion",
+        created=int(time.time()),
+        choices=[
+            Choice(
+                finish_reason="stop",
+                index=0,
+                message=ChatCompletionMessage(
+                    content="Test response",
+                    role="assistant",
+                ),
+            )
+        ],
+        usage=CompletionUsage(
+            completion_tokens=10,
+            prompt_tokens=20,
+            total_tokens=30,
+        ),
+    )
+
+
+@pytest.fixture
+def mock_embedding_response():
+    return CreateEmbeddingResponse(
+        data=[
+            Embedding(
+                embedding=[0.1, 0.2, 0.3],
+                index=0,
+                object="embedding",
+            )
+        ],
+        model="text-embedding-3-small",
+        object="list",
+        usage=Usage(
+            prompt_tokens=10,
+            total_tokens=10,
+        ),
+    )
+
+
+def test_basic_completion(mock_client, mock_openai_response):
+    with patch("openai.resources.chat.completions.Completions.create", return_value=mock_openai_response):
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+        response = client.chat.completions.create(
+            model="gpt-4",
+            messages=[{"role": "user", "content": "Hello"}],
+            posthog_distinct_id="test-id",
+            posthog_properties={"foo": "bar"},
+        )
+
+        assert response == mock_openai_response
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+
+        assert call_args["distinct_id"] == "test-id"
+        assert call_args["event"] == "$ai_generation"
+        assert props["$ai_provider"] == "openai"
+        assert props["$ai_model"] == "gpt-4"
+        assert props["$ai_input"] == [{"role": "user", "content": "Hello"}]
+        assert props["$ai_output"] == {"choices": [{"role": "assistant", "content": "Test response"}]}
+        assert props["$ai_input_tokens"] == 20
+        assert props["$ai_output_tokens"] == 10
+        assert props["$ai_http_status"] == 200
+        assert props["foo"] == "bar"
+        assert isinstance(props["$ai_latency"], float)
+
+
+def test_embeddings(mock_client, mock_embedding_response):
+    with patch("openai.resources.embeddings.Embeddings.create", return_value=mock_embedding_response):
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+        response = client.embeddings.create(
+            model="text-embedding-3-small",
+            input="Hello world",
+            posthog_distinct_id="test-id",
+            posthog_properties={"foo": "bar"},
+        )
+
+        assert response == mock_embedding_response
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+
+        assert call_args["distinct_id"] == "test-id"
+        assert call_args["event"] == "$ai_embedding"
+        assert props["$ai_provider"] == "openai"
+        assert props["$ai_model"] == "text-embedding-3-small"
+        assert props["$ai_input"] == "Hello world"
+        assert props["$ai_input_tokens"] == 10
+        assert props["$ai_http_status"] == 200
+        assert props["foo"] == "bar"
+        assert isinstance(props["$ai_latency"], float)
@@ -715,6 +715,25 @@ class TestClient(unittest.TestCase):
        self.assertTrue(isinstance(msg["timestamp"], str))
        self.assertIsNone(msg.get("uuid"))

+    def test_basic_group_identify_with_distinct_id(self):
+        success, msg = self.client.group_identify("organization", "id:5", distinct_id="distinct_id")
+        self.assertTrue(success)
+        self.assertEqual(msg["event"], "$groupidentify")
+        self.assertEqual(msg["distinct_id"], "distinct_id")
+        self.assertEqual(
+            msg["properties"],
+            {
+                "$group_type": "organization",
+                "$group_key": "id:5",
+                "$group_set": {},
+                "$lib": "posthog-python",
+                "$lib_version": VERSION,
+                "$geoip_disable": True,
+            },
+        )
+        self.assertTrue(isinstance(msg["timestamp"], str))
+        self.assertIsNone(msg.get("uuid"))
+
    def test_advanced_group_identify(self):
        success, msg = self.client.group_identify(
            "organization", "id:5", {"trait": "value"}, {"ip": "192.168.0.1"}, datetime(2014, 9, 3), "new-uuid"
@@ -737,6 +756,35 @@ class TestClient(unittest.TestCase):
        self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00")
        self.assertEqual(msg["context"]["ip"], "192.168.0.1")

+    def test_advanced_group_identify_with_distinct_id(self):
+        success, msg = self.client.group_identify(
+            "organization",
+            "id:5",
+            {"trait": "value"},
+            {"ip": "192.168.0.1"},
+            datetime(2014, 9, 3),
+            "new-uuid",
+            distinct_id="distinct_id",
+        )
+
+        self.assertTrue(success)
+        self.assertEqual(msg["event"], "$groupidentify")
+        self.assertEqual(msg["distinct_id"], "distinct_id")
+
+        self.assertEqual(
+            msg["properties"],
+            {
+                "$group_type": "organization",
+                "$group_key": "id:5",
+                "$group_set": {"trait": "value"},
+                "$lib": "posthog-python",
+                "$lib_version": VERSION,
+                "$geoip_disable": True,
+            },
+        )
+        self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00")
+        self.assertEqual(msg["context"]["ip"], "192.168.0.1")
+
    def test_basic_alias(self):
        client = self.client
        success, msg = client.alias("previousId", "distinct_id")
@@ -1,4 +1,4 @@
-VERSION = "3.7.4"
+VERSION = "3.8.2"

 if __name__ == "__main__":
    print(VERSION, end="")  # noqa: T201
@@ -1,2 +1,5 @@
 [bdist_wheel]
 universal = 1
+
+[tool:pytest]
+asyncio_mode = auto
@@ -14,7 +14,13 @@ long_description = """
 PostHog is developer-friendly, self-hosted product analytics. posthog-python is the python package.
 """

-install_requires = ["requests>=2.7,<3.0", "six>=1.5", "monotonic>=1.5", "backoff>=1.10.0", "python-dateutil>2.1"]
+install_requires = [
+    "requests>=2.7,<3.0",
+    "six>=1.5",
+    "monotonic>=1.5",
+    "backoff>=1.10.0",
+    "python-dateutil>2.1",
+]

 extras_require = {
    "dev": [
@@ -24,8 +30,21 @@ extras_require = {
        "flake8-print",
        "pre-commit",
    ],
-    "test": ["mock>=2.0.0", "freezegun==0.3.15", "pylint", "flake8", "coverage", "pytest", "pytest-timeout", "django"],
+    "test": [
+        "mock>=2.0.0",
+        "freezegun==0.3.15",
+        "pylint",
+        "flake8",
+        "coverage",
+        "pytest",
+        "pytest-timeout",
+        "pytest-asyncio",
+        "django",
+        "langchain-community>=0.2.0",
+        "langchain-openai>=0.2.0",
+    ],
    "sentry": ["sentry-sdk", "django"],
+    "langchain": ["langchain>=0.2.0"],
 }

 setup(
@@ -37,7 +56,15 @@ setup(
    maintainer="PostHog",
    maintainer_email="hey@posthog.com",
    test_suite="posthog.test.all",
-    packages=["posthog", "posthog.test", "posthog.sentry", "posthog.exception_integrations"],
+    packages=[
+        "posthog",
+        "posthog.ai",
+        "posthog.ai.langchain",
+        "posthog.ai.openai",
+        "posthog.test",
+        "posthog.sentry",
+        "posthog.exception_integrations",
+    ],
    license="MIT License",
    install_requires=install_requires,
    extras_require=extras_require,
@@ -60,5 +87,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",
    ],
 )
@@ -29,6 +29,7 @@ setup(
    test_suite="posthoganalytics.test.all",
    packages=[
        "posthoganalytics",
+        "posthoganalytics.ai",
        "posthoganalytics.test",
        "posthoganalytics.sentry",
        "posthoganalytics.exception_integrations",
Author	SHA1	Message	Date
Georgiy TarasovandGitHub	78ab0ca8b5	fix(llm-observability): include the ai packages (#162 ) * fix: setuptools * fix: include packages	2025-01-14 10:27:05 +01:00
Peter KirkhamandGitHub	c5bfc1377a	fix: update to export module (#161 )	2025-01-14 01:25:00 +00:00
Peter KirkhamandGitHub	6b1c0dc313	feat: Embeddings + Personless events + Destructure property JSON (#160 )	2025-01-14 00:35:09 +00:00
Georgiy TarasovandGitHub	e51b883e7b	feat(llm-observability): add langchain integration (#159 ) * feat(ai): LangChain integration v0.1 * test: langchain integration tests * test: langchain-openai for v2 and v3 * chore: reorganize imports * fix: ci * fix: set python on ci to 3.9 * fix: upgrade ci for python 3.9 * fix: fallback for distinct_id * fix: personless events for omitted distinct_ids * fix: review comments * feat: base url retrieval	2025-01-13 18:40:02 +01:00
Peter Kirkham GitHub Michael Matloka	66101c92bf	Feat: Add llm observability to python sdk (#158 ) Co-authored-by: Michael Matloka <michael@matloka.com>	2025-01-11 01:34:27 +00:00
Sibin M SandGitHub	05932b3f13	[FEATURE]Add distinct_id to group_identify (#155 ) * [FEATURE]Add distinct_id to group_identify * [TESTS]Updated test cases for adding distinct_id to group_identify * [LINT-FIX]client.py and test_client.py * [CHORE]Verion bump and changelog update	2025-01-03 16:00:35 -05:00
Dylan MartinandGitHub	50c13563b2	fix: CI (#156 ) * test CI * heck it, upgrade python * okay don't do anything silly with the cache hits i guess * more CI upgrades :crossedfingers * upgrade all CI to latest versions, then * jk this is how python works * whackamole * what even * yeesh * this can't be it * if this breaks ill kms * dark magic dark MAGIC * im giving up on my dreams	2025-01-03 15:50:21 -05:00
Dylan MartinandGitHub	dca4af66ae	Update CODEOWNERS (#154 )	2025-01-02 12:52:51 -05:00