feat: add support for responses api (#205 )

* feat: add suppoort for responses api * fix: test * fix: black * fix: test - hopefully * fix: test - hopefully #2 * fix: test - hopefully #3 * fix: test - hopefully #4 * fix: greptaile catch * fix: mypy is not my friend * fix: isort usort weallsort * fix: noredef * fix: mypy baseline * fix: mypy * fix: mypy
fix: distro requirement for analytics package (#204 )
2025-03-14 05:16:52 +00:00 · 2025-03-12 14:12:15 +00:00 · 2025-03-10 21:35:36 +00:00 · 2025-03-06 22:37:21 +00:00 · 2025-03-03 14:00:52 -05:00 · 2025-02-28 20:40:09 +00:00
24 changed files with 1596 additions and 333 deletions
@@ -13,10 +13,10 @@ jobs:
              with:
                  fetch-depth: 1

-            - name: Set up Python 3.13
+            - name: Set up Python 3.11
              uses: actions/setup-python@v2
              with:
-                  python-version: "3.13"
+                  python-version: 3.11.11

            - uses: actions/cache@v3
              with:
@@ -1,4 +1,50 @@

+
+## 3.20.0 – 2025-03-13
+
+1. Add support for OpenAI Responses API.
+
+## 3.19.2 – 2025-03-11
+
+1. Fix install requirements for analytics package
+
+## 3.19.1 – 2025-03-11
+
+1. Fix bug where None is sent as delta in azure
+
+## 3.19.0 – 2025-03-04
+
+1. Add support for tool calls in OpenAI and Anthropic.
+2. Add support for cached tokens.
+
+## 3.18.1 – 2025-03-03
+
+1. Improve quota-limited feature flag logs
+
+## 3.18.0 - 2025-02-28
+
+1. Add support for Azure OpenAI.
+
+## 3.17.0 - 2025-02-27
+
+1. The LangChain handler now captures tools in `$ai_generation` events, in property `$ai_tools`. This allows for displaying tools provided to the LLM call in PostHog UI. Note that support for `$ai_tools` in OpenAI and Anthropic SDKs is coming soon.
+
+## 3.16.0 - 2025-02-26
+
+1. feat: add some platform info to events (#198)
+
+## 3.15.1 - 2025-02-23
+
+1. Fix async client support for OpenAI.
+
+## 3.15.0 - 2025-02-19
+
+1. Support quota-limited feature flags
+
+## 3.14.2 - 2025-02-19
+
+1. Evaluate feature flag payloads with case sensitivity correctly. Fixes <https://github.com/PostHog/posthog-python/issues/178>
+
 ## 3.14.1 - 2025-02-18

 1. Add support for Bedrock Anthropic Usage
@@ -10,8 +10,10 @@ Please see the [Python integration docs](https://posthog.com/docs/integrations/p
 ### Testing Locally

 1. Run `python3 -m venv env` (creates virtual environment called "env")
+    * or `uv venv env`
 2. Run `source env/bin/activate` (activates the virtual environment)
 3. Run `python3 -m pip install -e ".[test]"` (installs the package in develop mode, along with test dependencies)
+    * or `uv pip install -e ".[test]"`
 4. Run `make test`
  1. To run a specific test do `pytest -k test_no_api_key`

@@ -1,220 +0,0 @@
-import os
-import uuid
-
-from pydantic import BaseModel
-
-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.host = os.getenv("POSTHOG_HOST", "http://localhost:8000")  # Or https://app.posthog.com
-posthog.debug = True
-# change this to False to see usage events
-# posthog.privacy_mode = True
-
-openai_client = OpenAI(
-    api_key=os.getenv("OPENAI_API_KEY", "your-openai-api-key"),
-    posthog_client=posthog,
-)
-
-async_openai_client = AsyncOpenAI(
-    api_key=os.getenv("OPENAI_API_KEY", "your-openai-api-key"),
-    posthog_client=posthog,
-)
-
-
-def main_sync():
-    trace_id = str(uuid.uuid4())
-    print("Trace ID:", trace_id)
-    distinct_id = "test2_distinct_id"
-    properties = {"test_property": "test_value"}
-    groups = {"company": "test_company"}
-
-    try:
-        # basic_openai_call(distinct_id, trace_id, properties, groups)
-        # streaming_openai_call(distinct_id, trace_id, properties, groups)
-        # embedding_openai_call(distinct_id, trace_id, properties, groups)
-        # image_openai_call()
-        beta_openai_call(distinct_id, trace_id, properties, groups)
-    except Exception as e:
-        print("Error during OpenAI call:", str(e))
-
-
-async def main_async():
-    trace_id = str(uuid.uuid4())
-    print("Trace ID:", trace_id)
-    distinct_id = "test_distinct_id"
-    properties = {"test_property": "test_value"}
-    groups = {"company": "test_company"}
-
-    try:
-        await basic_async_openai_call(distinct_id, trace_id, properties, groups)
-        await streaming_async_openai_call(distinct_id, trace_id, properties, groups)
-        await embedding_async_openai_call(distinct_id, trace_id, properties, groups)
-        await image_async_openai_call()
-    except Exception as e:
-        print("Error during OpenAI call:", str(e))
-
-
-def basic_openai_call(distinct_id, trace_id, properties, groups):
-    response = openai_client.chat.completions.create(
-        model="gpt-4o-mini",
-        messages=[
-            {"role": "system", "content": "You are a complex problem solver."},
-            {"role": "user", "content": "Explain quantum computing in simple terms."},
-        ],
-        max_tokens=100,
-        temperature=0.7,
-        posthog_distinct_id=distinct_id,
-        posthog_trace_id=trace_id,
-        posthog_properties=properties,
-        posthog_groups=groups,
-    )
-    print(response)
-    if response and response.choices:
-        print("OpenAI response:", response.choices[0].message.content)
-    else:
-        print("No response or unexpected format returned.")
-    return response
-
-
-async def basic_async_openai_call(distinct_id, trace_id, properties, groups):
-    response = await async_openai_client.chat.completions.create(
-        model="gpt-4o-mini",
-        messages=[
-            {"role": "system", "content": "You are a complex problem solver."},
-            {"role": "user", "content": "Explain quantum computing in simple terms."},
-        ],
-        max_tokens=100,
-        temperature=0.7,
-        posthog_distinct_id=distinct_id,
-        posthog_trace_id=trace_id,
-        posthog_properties=properties,
-        posthog_groups=groups,
-    )
-    if response and hasattr(response, "choices"):
-        print("OpenAI response:", response.choices[0].message.content)
-    else:
-        print("No response or unexpected format returned.")
-    return response
-
-
-def streaming_openai_call(distinct_id, trace_id, properties, groups):
-
-    response = openai_client.chat.completions.create(
-        model="gpt-4o-mini",
-        messages=[
-            {"role": "system", "content": "You are a complex problem solver."},
-            {"role": "user", "content": "Explain quantum computing in simple terms."},
-        ],
-        max_tokens=100,
-        temperature=0.7,
-        stream=True,
-        posthog_distinct_id=distinct_id,
-        posthog_trace_id=trace_id,
-        posthog_properties=properties,
-        posthog_groups=groups,
-    )
-
-    for chunk in response:
-        if hasattr(chunk, "choices") and chunk.choices and len(chunk.choices) > 0:
-            print(chunk.choices[0].delta.content or "", end="")
-
-    return response
-
-
-async def streaming_async_openai_call(distinct_id, trace_id, properties, groups):
-    response = await async_openai_client.chat.completions.create(
-        model="gpt-4o-mini",
-        messages=[
-            {"role": "system", "content": "You are a complex problem solver."},
-            {"role": "user", "content": "Explain quantum computing in simple terms."},
-        ],
-        max_tokens=100,
-        temperature=0.7,
-        stream=True,
-        posthog_distinct_id=distinct_id,
-        posthog_trace_id=trace_id,
-        posthog_properties=properties,
-        posthog_groups=groups,
-    )
-
-    async for chunk in response:
-        if hasattr(chunk, "choices") and chunk.choices and len(chunk.choices) > 0:
-            print(chunk.choices[0].delta.content or "", end="")
-
-    return response
-
-
-# none instrumented
-def image_openai_call():
-    response = openai_client.images.generate(model="dall-e-3", prompt="A cute baby hedgehog", n=1, size="1024x1024")
-    print(response)
-    return response
-
-
-# none instrumented
-async def image_async_openai_call():
-    response = await async_openai_client.images.generate(
-        model="dall-e-3", prompt="A cute baby hedgehog", n=1, size="1024x1024"
-    )
-    print(response)
-    return response
-
-
-def embedding_openai_call(posthog_distinct_id, posthog_trace_id, posthog_properties, posthog_groups):
-    response = openai_client.embeddings.create(
-        input="The hedgehog is cute",
-        model="text-embedding-3-small",
-        posthog_distinct_id=posthog_distinct_id,
-        posthog_trace_id=posthog_trace_id,
-        posthog_properties=posthog_properties,
-        posthog_groups=posthog_groups,
-    )
-    print(response)
-    return response
-
-
-async def embedding_async_openai_call(posthog_distinct_id, posthog_trace_id, posthog_properties, posthog_groups):
-    response = await async_openai_client.embeddings.create(
-        input="The hedgehog is cute",
-        model="text-embedding-3-small",
-        posthog_distinct_id=posthog_distinct_id,
-        posthog_trace_id=posthog_trace_id,
-        posthog_properties=posthog_properties,
-        posthog_groups=posthog_groups,
-    )
-    print(response)
-    return response
-
-
-class CalendarEvent(BaseModel):
-    name: str
-    date: str
-    participants: list[str]
-
-
-def beta_openai_call(distinct_id, trace_id, properties, groups):
-    response = openai_client.beta.chat.completions.parse(
-        model="gpt-4o-mini",
-        messages=[
-            {"role": "system", "content": "Extract the event information."},
-            {"role": "user", "content": "Alice and Bob are going to a science fair on Friday."},
-        ],
-        response_format=CalendarEvent,
-        posthog_distinct_id=distinct_id,
-        posthog_trace_id=trace_id,
-        posthog_properties=properties,
-        posthog_groups=groups,
-    )
-    print(response)
-    return response
-
-
-# HOW TO RUN:
-# comment out one of these to run the other
-
-if __name__ == "__main__":
-    main_sync()
-# asyncio.run(main_async())
@@ -35,27 +35,7 @@ posthog/sentry/posthog_integration.py:0: error: Statement is unreachable  [unrea
 posthog/ai/utils.py:0: error: Need type annotation for "output" (hint: "output: list[<type>] = ...")  [var-annotated]
 posthog/ai/utils.py:0: error: Function "builtins.any" is not valid as a type  [valid-type]
 posthog/ai/utils.py:0: note: Perhaps you meant "typing.Any" instead of "any"?
-posthog/ai/utils.py:0: error: Incompatible types in assignment (expression has type "UUID", variable has type "str | None")  [assignment]
 posthog/ai/utils.py:0: error: Function "builtins.any" is not valid as a type  [valid-type]
 posthog/ai/utils.py:0: note: Perhaps you meant "typing.Any" instead of "any"?
-posthog/ai/utils.py:0: error: Incompatible types in assignment (expression has type "UUID", variable has type "str | None")  [assignment]
 sentry_django_example/sentry_django_example/settings.py:0: error: Need type annotation for "ALLOWED_HOSTS" (hint: "ALLOWED_HOSTS: list[<type>] = ...")  [var-annotated]
-sentry_django_example/sentry_django_example/settings.py:0: error: Incompatible types in assignment (expression has type "str", variable has type "None")  [assignment]
-posthog/ai/openai/openai_async.py:0: error: Incompatible types in assignment (expression has type "UUID", variable has type "str | None")  [assignment]
-posthog/ai/openai/openai_async.py:0: error: Incompatible types in assignment (expression has type "UUID", variable has type "str | None")  [assignment]
-posthog/ai/openai/openai_async.py:0: error: Unpacked dict entry 11 has incompatible type "dict[str, Any] | None"; expected "SupportsKeysAndGetItem[str, Any]"  [dict-item]
-posthog/ai/openai/openai_async.py:0: error: Incompatible types in assignment (expression has type "UUID", variable has type "str | None")  [assignment]
-posthog/ai/openai/openai_async.py:0: error: Unpacked dict entry 8 has incompatible type "dict[str, Any] | None"; expected "SupportsKeysAndGetItem[str, Any]"  [dict-item]
-posthog/ai/openai/openai.py:0: error: Incompatible types in assignment (expression has type "UUID", variable has type "str | None")  [assignment]
-posthog/ai/openai/openai.py:0: error: Incompatible types in assignment (expression has type "UUID", variable has type "str | None")  [assignment]
-posthog/ai/openai/openai.py:0: error: Unpacked dict entry 11 has incompatible type "dict[str, Any] | None"; expected "SupportsKeysAndGetItem[str, Any]"  [dict-item]
-posthog/ai/openai/openai.py:0: error: Incompatible types in assignment (expression has type "UUID", variable has type "str | None")  [assignment]
-posthog/ai/openai/openai.py:0: error: Unpacked dict entry 8 has incompatible type "dict[str, Any] | None"; expected "SupportsKeysAndGetItem[str, Any]"  [dict-item]
-posthog/ai/anthropic/anthropic_async.py:0: error: Incompatible types in assignment (expression has type "UUID", variable has type "str | None")  [assignment]
-posthog/ai/anthropic/anthropic_async.py:0: error: Incompatible types in assignment (expression has type "UUID", variable has type "str | None")  [assignment]
-posthog/ai/anthropic/anthropic_async.py:0: error: Incompatible types in assignment (expression has type "UUID", variable has type "str | None")  [assignment]
-posthog/ai/anthropic/anthropic.py:0: error: Incompatible types in assignment (expression has type "UUID", variable has type "str | None")  [assignment]
-posthog/ai/anthropic/anthropic.py:0: error: Incompatible types in assignment (expression has type "UUID", variable has type "str | None")  [assignment]
-posthog/ai/anthropic/anthropic.py:0: error: Incompatible types in assignment (expression has type "UUID", variable has type "str | None")  [assignment]
-llm_observability_examples.py:0: error: Argument "posthog_client" to "OpenAI" has incompatible type Module; expected "Client"  [arg-type]
-llm_observability_examples.py:0: error: Argument "posthog_client" to "AsyncOpenAI" has incompatible type Module; expected "Client"  [arg-type]
+sentry_django_example/sentry_django_example/settings.py:0: error: Incompatible types in assignment (expression has type "str", variable has type "None")  [assignment]
@@ -54,7 +54,7 @@ class WrappedMessages(Messages):
            **kwargs: Arguments passed to Anthropic's messages.create
        """
        if posthog_trace_id is None:
-            posthog_trace_id = uuid.uuid4()
+            posthog_trace_id = str(uuid.uuid4())

        if kwargs.get("stream", False):
            return self._create_streaming(
@@ -89,7 +89,7 @@ class WrappedMessages(Messages):
        **kwargs: Any,
    ):
        if posthog_trace_id is None:
-            posthog_trace_id = uuid.uuid4()
+            posthog_trace_id = str(uuid.uuid4())

        return self._create_streaming(
            posthog_distinct_id,
@@ -125,6 +125,8 @@ class WrappedMessages(Messages):
                            for k in [
                                "input_tokens",
                                "output_tokens",
+                                "cache_read_input_tokens",
+                                "cache_creation_input_tokens",
                            ]
                        }

@@ -165,7 +167,7 @@ class WrappedMessages(Messages):
        output: str,
    ):
        if posthog_trace_id is None:
-            posthog_trace_id = uuid.uuid4()
+            posthog_trace_id = str(uuid.uuid4())

        event_properties = {
            "$ai_provider": "anthropic",
@@ -184,6 +186,8 @@ class WrappedMessages(Messages):
            "$ai_http_status": 200,
            "$ai_input_tokens": usage_stats.get("input_tokens", 0),
            "$ai_output_tokens": usage_stats.get("output_tokens", 0),
+            "$ai_cache_read_input_tokens": usage_stats.get("cache_read_input_tokens", 0),
+            "$ai_cache_creation_input_tokens": usage_stats.get("cache_creation_input_tokens", 0),
            "$ai_latency": latency,
            "$ai_trace_id": posthog_trace_id,
            "$ai_base_url": str(self._client.base_url),
@@ -54,7 +54,7 @@ class AsyncWrappedMessages(AsyncMessages):
            **kwargs: Arguments passed to Anthropic's messages.create
        """
        if posthog_trace_id is None:
-            posthog_trace_id = uuid.uuid4()
+            posthog_trace_id = str(uuid.uuid4())

        if kwargs.get("stream", False):
            return await self._create_streaming(
@@ -89,7 +89,7 @@ class AsyncWrappedMessages(AsyncMessages):
        **kwargs: Any,
    ):
        if posthog_trace_id is None:
-            posthog_trace_id = uuid.uuid4()
+            posthog_trace_id = str(uuid.uuid4())

        return await self._create_streaming(
            posthog_distinct_id,
@@ -125,6 +125,8 @@ class AsyncWrappedMessages(AsyncMessages):
                            for k in [
                                "input_tokens",
                                "output_tokens",
+                                "cache_read_input_tokens",
+                                "cache_creation_input_tokens",
                            ]
                        }

@@ -165,7 +167,7 @@ class AsyncWrappedMessages(AsyncMessages):
        output: str,
    ):
        if posthog_trace_id is None:
-            posthog_trace_id = uuid.uuid4()
+            posthog_trace_id = str(uuid.uuid4())

        event_properties = {
            "$ai_provider": "anthropic",
@@ -184,6 +186,8 @@ class AsyncWrappedMessages(AsyncMessages):
            "$ai_http_status": 200,
            "$ai_input_tokens": usage_stats.get("input_tokens", 0),
            "$ai_output_tokens": usage_stats.get("output_tokens", 0),
+            "$ai_cache_read_input_tokens": usage_stats.get("cache_read_input_tokens", 0),
+            "$ai_cache_creation_input_tokens": usage_stats.get("cache_creation_input_tokens", 0),
            "$ai_latency": latency,
            "$ai_trace_id": posthog_trace_id,
            "$ai_base_url": str(self._client.base_url),
@@ -60,6 +60,8 @@ class GenerationMetadata(SpanMetadata):
    """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."""
+    tools: Optional[List[Dict[str, Any]]] = None
+    """Tools provided to the model."""


 RunMetadata = Union[SpanMetadata, GenerationMetadata]
@@ -377,6 +379,8 @@ class CallbackHandler(BaseCallbackHandler):
        generation = GenerationMetadata(name=run_name, input=messages, start_time=time.time(), end_time=None)
        if isinstance(invocation_params, dict):
            generation.model_params = get_model_params(invocation_params)
+            if tools := invocation_params.get("tools"):
+                generation.tools = tools
        if isinstance(metadata, dict):
            if model := metadata.get("ls_model_name"):
                generation.model = model
@@ -424,7 +428,11 @@ class CallbackHandler(BaseCallbackHandler):
            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)
+            trace_id,
+            run_id,
+            run,
+            outputs,
+            self._get_parent_run_id(trace_id, run_id, parent_run_id),
        )

    def _capture_trace_or_span(
@@ -465,7 +473,10 @@ class CallbackHandler(BaseCallbackHandler):
        )

    def _pop_run_and_capture_generation(
-        self, run_id: UUID, parent_run_id: Optional[UUID], response: Union[LLMResult, BaseException]
+        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)
@@ -476,7 +487,11 @@ class CallbackHandler(BaseCallbackHandler):
            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)
+            trace_id,
+            run_id,
+            run,
+            response,
+            self._get_parent_run_id(trace_id, run_id, parent_run_id),
        )

    def _capture_generation(
@@ -500,6 +515,12 @@ class CallbackHandler(BaseCallbackHandler):
            "$ai_latency": run.latency,
            "$ai_base_url": run.base_url,
        }
+        if run.tools:
+            event_properties["$ai_tools"] = with_privacy_mode(
+                self._client,
+                self._privacy_mode,
+                run.tools,
+            )

        if isinstance(output, BaseException):
            event_properties["$ai_http_status"] = _get_http_status(output)
@@ -1,4 +1,5 @@
 from .openai import OpenAI
 from .openai_async import AsyncOpenAI
+from .openai_providers import AsyncAzureOpenAI, AzureOpenAI

-__all__ = ["OpenAI", "AsyncOpenAI"]
+__all__ = ["OpenAI", "AsyncOpenAI", "AzureOpenAI", "AsyncAzureOpenAI"]
@@ -1,6 +1,6 @@
 import time
 import uuid
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional

 try:
    import openai
@@ -32,6 +32,167 @@ class OpenAI(openai.OpenAI):
        self.chat = WrappedChat(self)
        self.embeddings = WrappedEmbeddings(self)
        self.beta = WrappedBeta(self)
+        self.responses = WrappedResponses(self)
+
+
+class WrappedResponses(openai.resources.responses.Responses):
+    _client: OpenAI
+
+    def create(
+        self,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: bool = False,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        if posthog_trace_id is None:
+            posthog_trace_id = str(uuid.uuid4())
+
+        if kwargs.get("stream", False):
+            return self._create_streaming(
+                posthog_distinct_id,
+                posthog_trace_id,
+                posthog_properties,
+                posthog_privacy_mode,
+                posthog_groups,
+                **kwargs,
+            )
+
+        return call_llm_and_track_usage(
+            posthog_distinct_id,
+            self._client._ph_client,
+            "openai",
+            posthog_trace_id,
+            posthog_properties,
+            posthog_privacy_mode,
+            posthog_groups,
+            self._client.base_url,
+            super().create,
+            **kwargs,
+        )
+
+    def _create_streaming(
+        self,
+        posthog_distinct_id: Optional[str],
+        posthog_trace_id: Optional[str],
+        posthog_properties: Optional[Dict[str, Any]],
+        posthog_privacy_mode: bool,
+        posthog_groups: Optional[Dict[str, Any]],
+        **kwargs: Any,
+    ):
+        start_time = time.time()
+        usage_stats: Dict[str, int] = {}
+        final_content = []
+        response = super().create(**kwargs)
+
+        def generator():
+            nonlocal usage_stats
+            nonlocal final_content
+
+            try:
+                for chunk in response:
+                    if hasattr(chunk, "type") and chunk.type == "response.completed":
+                        res = chunk.response
+                        if res.output and len(res.output) > 0:
+                            final_content.append(res.output[0])
+
+                    if hasattr(chunk, "usage") and chunk.usage:
+                        usage_stats = {
+                            k: getattr(chunk.usage, k, 0)
+                            for k in [
+                                "input_tokens",
+                                "output_tokens",
+                                "total_tokens",
+                            ]
+                        }
+
+                        # Add support for cached tokens
+                        if hasattr(chunk.usage, "output_tokens_details") and hasattr(
+                            chunk.usage.output_tokens_details, "reasoning_tokens"
+                        ):
+                            usage_stats["reasoning_tokens"] = chunk.usage.output_tokens_details.reasoning_tokens
+
+                        if hasattr(chunk.usage, "input_tokens_details") and hasattr(
+                            chunk.usage.input_tokens_details, "cached_tokens"
+                        ):
+                            usage_stats["cache_read_input_tokens"] = chunk.usage.input_tokens_details.cached_tokens
+
+                    yield chunk
+
+            finally:
+                end_time = time.time()
+                latency = end_time - start_time
+                output = final_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: Any,
+        tool_calls: Optional[List[Dict[str, Any]]] = None,
+    ):
+        if posthog_trace_id is None:
+            posthog_trace_id = str(uuid.uuid4())
+
+        event_properties = {
+            "$ai_provider": "openai",
+            "$ai_model": kwargs.get("model"),
+            "$ai_model_parameters": get_model_params(kwargs),
+            "$ai_input": with_privacy_mode(self._client._ph_client, posthog_privacy_mode, kwargs.get("input")),
+            "$ai_output_choices": with_privacy_mode(
+                self._client._ph_client,
+                posthog_privacy_mode,
+                output,
+            ),
+            "$ai_http_status": 200,
+            "$ai_input_tokens": usage_stats.get("input_tokens", 0),
+            "$ai_output_tokens": usage_stats.get("output_tokens", 0),
+            "$ai_cache_read_input_tokens": usage_stats.get("cache_read_input_tokens", 0),
+            "$ai_reasoning_tokens": usage_stats.get("reasoning_tokens", 0),
+            "$ai_latency": latency,
+            "$ai_trace_id": posthog_trace_id,
+            "$ai_base_url": str(self._client.base_url),
+            **(posthog_properties or {}),
+        }
+
+        if tool_calls:
+            event_properties["$ai_tools"] = with_privacy_mode(
+                self._client._ph_client,
+                posthog_privacy_mode,
+                tool_calls,
+            )
+
+        if posthog_distinct_id is None:
+            event_properties["$process_person_profile"] = False
+
+        if hasattr(self._client._ph_client, "capture"):
+            self._client._ph_client.capture(
+                distinct_id=posthog_distinct_id or posthog_trace_id,
+                event="$ai_generation",
+                properties=event_properties,
+                groups=posthog_groups,
+            )


 class WrappedChat(openai.resources.chat.Chat):
@@ -55,7 +216,7 @@ class WrappedCompletions(openai.resources.chat.completions.Completions):
        **kwargs: Any,
    ):
        if posthog_trace_id is None:
-            posthog_trace_id = uuid.uuid4()
+            posthog_trace_id = str(uuid.uuid4())

        if kwargs.get("stream", False):
            return self._create_streaming(
@@ -92,6 +253,7 @@ class WrappedCompletions(openai.resources.chat.completions.Completions):
        start_time = time.time()
        usage_stats: Dict[str, int] = {}
        accumulated_content = []
+        accumulated_tools = {}
        if "stream_options" not in kwargs:
            kwargs["stream_options"] = {}
        kwargs["stream_options"]["include_usage"] = True
@@ -100,6 +262,8 @@ class WrappedCompletions(openai.resources.chat.completions.Completions):
        def generator():
            nonlocal usage_stats
            nonlocal accumulated_content
+            nonlocal accumulated_tools
+
            try:
                for chunk in response:
                    if hasattr(chunk, "usage") and chunk.usage:
@@ -112,10 +276,34 @@ class WrappedCompletions(openai.resources.chat.completions.Completions):
                            ]
                        }

+                        # Add support for cached tokens
+                        if hasattr(chunk.usage, "prompt_tokens_details") and hasattr(
+                            chunk.usage.prompt_tokens_details, "cached_tokens"
+                        ):
+                            usage_stats["cache_read_input_tokens"] = chunk.usage.prompt_tokens_details.cached_tokens
+
+                        if hasattr(chunk.usage, "output_tokens_details") and hasattr(
+                            chunk.usage.output_tokens_details, "reasoning_tokens"
+                        ):
+                            usage_stats["reasoning_tokens"] = chunk.usage.output_tokens_details.reasoning_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)
+                        if chunk.choices[0].delta and chunk.choices[0].delta.content:
+                            content = chunk.choices[0].delta.content
+                            if content:
+                                accumulated_content.append(content)
+
+                        # Process tool calls
+                        tool_calls = getattr(chunk.choices[0].delta, "tool_calls", None)
+                        if tool_calls:
+                            for tool_call in tool_calls:
+                                index = tool_call.index
+                                if index not in accumulated_tools:
+                                    accumulated_tools[index] = tool_call
+                                else:
+                                    # Append arguments for existing tool calls
+                                    if hasattr(tool_call, "function") and hasattr(tool_call.function, "arguments"):
+                                        accumulated_tools[index].function.arguments += tool_call.function.arguments

                    yield chunk

@@ -123,6 +311,7 @@ class WrappedCompletions(openai.resources.chat.completions.Completions):
                end_time = time.time()
                latency = end_time - start_time
                output = "".join(accumulated_content)
+                tools = list(accumulated_tools.values()) if accumulated_tools else None
                self._capture_streaming_event(
                    posthog_distinct_id,
                    posthog_trace_id,
@@ -133,6 +322,7 @@ class WrappedCompletions(openai.resources.chat.completions.Completions):
                    usage_stats,
                    latency,
                    output,
+                    tools,
                )

        return generator()
@@ -147,10 +337,11 @@ class WrappedCompletions(openai.resources.chat.completions.Completions):
        kwargs: Dict[str, Any],
        usage_stats: Dict[str, int],
        latency: float,
-        output: str,
+        output: Any,
+        tool_calls: Optional[List[Dict[str, Any]]] = None,
    ):
        if posthog_trace_id is None:
-            posthog_trace_id = uuid.uuid4()
+            posthog_trace_id = str(uuid.uuid4())

        event_properties = {
            "$ai_provider": "openai",
@@ -165,12 +356,21 @@ class WrappedCompletions(openai.resources.chat.completions.Completions):
            "$ai_http_status": 200,
            "$ai_input_tokens": usage_stats.get("prompt_tokens", 0),
            "$ai_output_tokens": usage_stats.get("completion_tokens", 0),
+            "$ai_cache_read_input_tokens": usage_stats.get("cache_read_input_tokens", 0),
+            "$ai_reasoning_tokens": usage_stats.get("reasoning_tokens", 0),
            "$ai_latency": latency,
            "$ai_trace_id": posthog_trace_id,
            "$ai_base_url": str(self._client.base_url),
-            **posthog_properties,
+            **(posthog_properties or {}),
        }

+        if tool_calls:
+            event_properties["$ai_tools"] = with_privacy_mode(
+                self._client._ph_client,
+                posthog_privacy_mode,
+                tool_calls,
+            )
+
        if posthog_distinct_id is None:
            event_properties["$process_person_profile"] = False

@@ -208,7 +408,7 @@ class WrappedEmbeddings(openai.resources.embeddings.Embeddings):
            The response from OpenAI's embeddings.create call.
        """
        if posthog_trace_id is None:
-            posthog_trace_id = uuid.uuid4()
+            posthog_trace_id = str(uuid.uuid4())

        start_time = time.time()
        response = super().create(**kwargs)
@@ -234,7 +434,7 @@ class WrappedEmbeddings(openai.resources.embeddings.Embeddings):
            "$ai_latency": latency,
            "$ai_trace_id": posthog_trace_id,
            "$ai_base_url": str(self._client.base_url),
-            **posthog_properties,
+            **(posthog_properties or {}),
        }

        if posthog_distinct_id is None:
@@ -1,6 +1,6 @@
 import time
 import uuid
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional

 try:
    import openai
@@ -31,17 +31,10 @@ class AsyncOpenAI(openai.AsyncOpenAI):
        self.chat = WrappedChat(self)
        self.embeddings = WrappedEmbeddings(self)
        self.beta = WrappedBeta(self)
+        self.responses = WrappedResponses(self)


-class WrappedChat(openai.resources.chat.AsyncChat):
-    _client: AsyncOpenAI
-
-    @property
-    def completions(self):
-        return WrappedCompletions(self._client)
-
-
-class WrappedCompletions(openai.resources.chat.completions.AsyncCompletions):
+class WrappedResponses(openai.resources.responses.Responses):
    _client: AsyncOpenAI

    async def create(
@@ -54,9 +47,8 @@ class WrappedCompletions(openai.resources.chat.completions.AsyncCompletions):
        **kwargs: Any,
    ):
        if posthog_trace_id is None:
-            posthog_trace_id = uuid.uuid4()
+            posthog_trace_id = str(uuid.uuid4())

-        # If streaming, handle streaming specifically
        if kwargs.get("stream", False):
            return await self._create_streaming(
                posthog_distinct_id,
@@ -67,59 +59,71 @@ class WrappedCompletions(openai.resources.chat.completions.AsyncCompletions):
                **kwargs,
            )

-        response = await call_llm_and_track_usage_async(
+        return await call_llm_and_track_usage_async(
            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,
        )
-        return response

    async def _create_streaming(
        self,
        posthog_distinct_id: Optional[str],
        posthog_trace_id: Optional[str],
        posthog_properties: Optional[Dict[str, Any]],
-        posthog_privacy_mode: bool = False,
-        posthog_groups: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: bool,
+        posthog_groups: Optional[Dict[str, Any]],
        **kwargs: Any,
    ):
        start_time = time.time()
        usage_stats: Dict[str, int] = {}
-        accumulated_content = []
-        if "stream_options" not in kwargs:
-            kwargs["stream_options"] = {}
-        kwargs["stream_options"]["include_usage"] = True
+        final_content = []
        response = await super().create(**kwargs)

        async def async_generator():
-            nonlocal usage_stats, accumulated_content
+            nonlocal usage_stats
+            nonlocal final_content
+
            try:
                async for chunk in response:
+                    if hasattr(chunk, "type") and chunk.type == "response.completed":
+                        res = chunk.response
+                        if res.output and len(res.output) > 0:
+                            final_content.append(res.output[0])
+
                    if hasattr(chunk, "usage") and chunk.usage:
                        usage_stats = {
                            k: getattr(chunk.usage, k, 0)
                            for k in [
-                                "prompt_tokens",
-                                "completion_tokens",
+                                "input_tokens",
+                                "output_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)
+
+                        # Add support for cached tokens
+                        if hasattr(chunk.usage, "output_tokens_details") and hasattr(
+                            chunk.usage.output_tokens_details, "reasoning_tokens"
+                        ):
+                            usage_stats["reasoning_tokens"] = chunk.usage.output_tokens_details.reasoning_tokens
+
+                        if hasattr(chunk.usage, "input_tokens_details") and hasattr(
+                            chunk.usage.input_tokens_details, "cached_tokens"
+                        ):
+                            usage_stats["cache_read_input_tokens"] = chunk.usage.input_tokens_details.cached_tokens

                    yield chunk

            finally:
                end_time = time.time()
                latency = end_time - start_time
-                output = "".join(accumulated_content)
+                output = final_content
                await self._capture_streaming_event(
                    posthog_distinct_id,
                    posthog_trace_id,
@@ -144,10 +148,193 @@ class WrappedCompletions(openai.resources.chat.completions.AsyncCompletions):
        kwargs: Dict[str, Any],
        usage_stats: Dict[str, int],
        latency: float,
-        output: str,
+        output: Any,
+        tool_calls: Optional[List[Dict[str, Any]]] = None,
    ):
        if posthog_trace_id is None:
-            posthog_trace_id = uuid.uuid4()
+            posthog_trace_id = str(uuid.uuid4())
+
+        event_properties = {
+            "$ai_provider": "openai",
+            "$ai_model": kwargs.get("model"),
+            "$ai_model_parameters": get_model_params(kwargs),
+            "$ai_input": with_privacy_mode(self._client._ph_client, posthog_privacy_mode, kwargs.get("input")),
+            "$ai_output_choices": with_privacy_mode(
+                self._client._ph_client,
+                posthog_privacy_mode,
+                output,
+            ),
+            "$ai_http_status": 200,
+            "$ai_input_tokens": usage_stats.get("input_tokens", 0),
+            "$ai_output_tokens": usage_stats.get("output_tokens", 0),
+            "$ai_cache_read_input_tokens": usage_stats.get("cache_read_input_tokens", 0),
+            "$ai_reasoning_tokens": usage_stats.get("reasoning_tokens", 0),
+            "$ai_latency": latency,
+            "$ai_trace_id": posthog_trace_id,
+            "$ai_base_url": str(self._client.base_url),
+            **(posthog_properties or {}),
+        }
+
+        if tool_calls:
+            event_properties["$ai_tools"] = with_privacy_mode(
+                self._client._ph_client,
+                posthog_privacy_mode,
+                tool_calls,
+            )
+
+        if posthog_distinct_id is None:
+            event_properties["$process_person_profile"] = False
+
+        if hasattr(self._client._ph_client, "capture"):
+            await self._client._ph_client.capture(
+                distinct_id=posthog_distinct_id or posthog_trace_id,
+                event="$ai_generation",
+                properties=event_properties,
+                groups=posthog_groups,
+            )
+
+
+class WrappedChat(openai.resources.chat.AsyncChat):
+    _client: AsyncOpenAI
+
+    @property
+    def completions(self):
+        return WrappedCompletions(self._client)
+
+
+class WrappedCompletions(openai.resources.chat.completions.AsyncCompletions):
+    _client: AsyncOpenAI
+
+    async def create(
+        self,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: bool = False,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        if posthog_trace_id is None:
+            posthog_trace_id = str(uuid.uuid4())
+
+        # If streaming, handle streaming specifically
+        if kwargs.get("stream", False):
+            return await self._create_streaming(
+                posthog_distinct_id,
+                posthog_trace_id,
+                posthog_properties,
+                posthog_privacy_mode,
+                posthog_groups,
+                **kwargs,
+            )
+
+        response = await call_llm_and_track_usage_async(
+            posthog_distinct_id,
+            self._client._ph_client,
+            "openai",
+            posthog_trace_id,
+            posthog_properties,
+            posthog_privacy_mode,
+            posthog_groups,
+            self._client.base_url,
+            super().create,
+            **kwargs,
+        )
+        return response
+
+    async def _create_streaming(
+        self,
+        posthog_distinct_id: Optional[str],
+        posthog_trace_id: Optional[str],
+        posthog_properties: Optional[Dict[str, Any]],
+        posthog_privacy_mode: bool = False,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        start_time = time.time()
+        usage_stats: Dict[str, int] = {}
+        accumulated_content = []
+        accumulated_tools = {}
+        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, accumulated_tools
+            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",
+                            ]
+                        }
+
+                        # Add support for cached tokens
+                        if hasattr(chunk.usage, "prompt_tokens_details") and hasattr(
+                            chunk.usage.prompt_tokens_details, "cached_tokens"
+                        ):
+                            usage_stats["cache_read_input_tokens"] = chunk.usage.prompt_tokens_details.cached_tokens
+
+                    if hasattr(chunk, "choices") and chunk.choices and len(chunk.choices) > 0:
+                        if chunk.choices[0].delta and chunk.choices[0].delta.content:
+                            content = chunk.choices[0].delta.content
+                            if content:
+                                accumulated_content.append(content)
+
+                        # Process tool calls
+                        tool_calls = getattr(chunk.choices[0].delta, "tool_calls", None)
+                        if tool_calls:
+                            for tool_call in tool_calls:
+                                index = tool_call.index
+                                if index not in accumulated_tools:
+                                    accumulated_tools[index] = tool_call
+                                else:
+                                    # Append arguments for existing tool calls
+                                    if hasattr(tool_call, "function") and hasattr(tool_call.function, "arguments"):
+                                        accumulated_tools[index].function.arguments += tool_call.function.arguments
+
+                    yield chunk
+
+            finally:
+                end_time = time.time()
+                latency = end_time - start_time
+                output = "".join(accumulated_content)
+                tools = list(accumulated_tools.values()) if accumulated_tools else None
+                await self._capture_streaming_event(
+                    posthog_distinct_id,
+                    posthog_trace_id,
+                    posthog_properties,
+                    posthog_privacy_mode,
+                    posthog_groups,
+                    kwargs,
+                    usage_stats,
+                    latency,
+                    output,
+                    tools,
+                )
+
+        return async_generator()
+
+    async def _capture_streaming_event(
+        self,
+        posthog_distinct_id: Optional[str],
+        posthog_trace_id: Optional[str],
+        posthog_properties: Optional[Dict[str, Any]],
+        posthog_privacy_mode: bool,
+        posthog_groups: Optional[Dict[str, Any]],
+        kwargs: Dict[str, Any],
+        usage_stats: Dict[str, int],
+        latency: float,
+        output: Any,
+        tool_calls: Optional[List[Dict[str, Any]]] = None,
+    ):
+        if posthog_trace_id is None:
+            posthog_trace_id = str(uuid.uuid4())

        event_properties = {
            "$ai_provider": "openai",
@@ -162,17 +349,25 @@ class WrappedCompletions(openai.resources.chat.completions.AsyncCompletions):
            "$ai_http_status": 200,
            "$ai_input_tokens": usage_stats.get("prompt_tokens", 0),
            "$ai_output_tokens": usage_stats.get("completion_tokens", 0),
+            "$ai_cache_read_input_tokens": usage_stats.get("cache_read_input_tokens", 0),
            "$ai_latency": latency,
            "$ai_trace_id": posthog_trace_id,
            "$ai_base_url": str(self._client.base_url),
-            **posthog_properties,
+            **(posthog_properties or {}),
        }

+        if tool_calls:
+            event_properties["$ai_tools"] = with_privacy_mode(
+                self._client._ph_client,
+                posthog_privacy_mode,
+                tool_calls,
+            )
+
        if posthog_distinct_id is None:
            event_properties["$process_person_profile"] = False

        if hasattr(self._client._ph_client, "capture"):
-            self._client._ph_client.capture(
+            await self._client._ph_client.capture(
                distinct_id=posthog_distinct_id or posthog_trace_id,
                event="$ai_generation",
                properties=event_properties,
@@ -207,7 +402,7 @@ class WrappedEmbeddings(openai.resources.embeddings.AsyncEmbeddings):
            The response from OpenAI's embeddings.create call.
        """
        if posthog_trace_id is None:
-            posthog_trace_id = uuid.uuid4()
+            posthog_trace_id = str(uuid.uuid4())

        start_time = time.time()
        response = await super().create(**kwargs)
@@ -233,7 +428,7 @@ class WrappedEmbeddings(openai.resources.embeddings.AsyncEmbeddings):
            "$ai_latency": latency,
            "$ai_trace_id": posthog_trace_id,
            "$ai_base_url": str(self._client.base_url),
-            **posthog_properties,
+            **(posthog_properties or {}),
        }

        if posthog_distinct_id is None:
@@ -0,0 +1,41 @@
+try:
+    import openai
+    import openai.resources
+except ImportError:
+    raise ModuleNotFoundError("Please install the Open AI SDK to use this feature: 'pip install openai'")
+
+from posthog.ai.openai.openai import WrappedBeta, WrappedChat, WrappedEmbeddings
+from posthog.ai.openai.openai_async import WrappedBeta as AsyncWrappedBeta
+from posthog.ai.openai.openai_async import WrappedChat as AsyncWrappedChat
+from posthog.ai.openai.openai_async import WrappedEmbeddings as AsyncWrappedEmbeddings
+from posthog.client import Client as PostHogClient
+
+
+class AzureOpenAI(openai.AzureOpenAI):
+    """
+    A wrapper around the Azure OpenAI 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.chat = WrappedChat(self)
+        self.embeddings = WrappedEmbeddings(self)
+        self.beta = WrappedBeta(self)
+
+
+class AsyncAzureOpenAI(openai.AsyncAzureOpenAI):
+    """
+    A wrapper around the Azure OpenAI 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.chat = AsyncWrappedChat(self)
+        self.embeddings = AsyncWrappedEmbeddings(self)
+        self.beta = AsyncWrappedBeta(self)
@@ -1,6 +1,6 @@
 import time
 import uuid
-from typing import Any, Callable, Dict, Optional
+from typing import Any, Callable, Dict, List, Optional

 from httpx import URL

@@ -34,15 +34,51 @@ def get_usage(response, provider: str) -> Dict[str, Any]:
        return {
            "input_tokens": response.usage.input_tokens,
            "output_tokens": response.usage.output_tokens,
+            "cache_read_input_tokens": response.usage.cache_read_input_tokens,
+            "cache_creation_input_tokens": response.usage.cache_creation_input_tokens,
        }
    elif provider == "openai":
+        cached_tokens = 0
+        input_tokens = 0
+        output_tokens = 0
+        reasoning_tokens = 0
+
+        # responses api
+        if hasattr(response.usage, "input_tokens"):
+            input_tokens = response.usage.input_tokens
+        if hasattr(response.usage, "output_tokens"):
+            output_tokens = response.usage.output_tokens
+        if hasattr(response.usage, "input_tokens_details") and hasattr(
+            response.usage.input_tokens_details, "cached_tokens"
+        ):
+            cached_tokens = response.usage.input_tokens_details.cached_tokens
+        if hasattr(response.usage, "output_tokens_details") and hasattr(
+            response.usage.output_tokens_details, "reasoning_tokens"
+        ):
+            reasoning_tokens = response.usage.output_tokens_details.reasoning_tokens
+
+        # chat completions
+        if hasattr(response.usage, "prompt_tokens"):
+            input_tokens = response.usage.prompt_tokens
+        if hasattr(response.usage, "completion_tokens"):
+            output_tokens = response.usage.completion_tokens
+        if hasattr(response.usage, "prompt_tokens_details") and hasattr(
+            response.usage.prompt_tokens_details, "cached_tokens"
+        ):
+            cached_tokens = response.usage.prompt_tokens_details.cached_tokens
+
        return {
-            "input_tokens": response.usage.prompt_tokens,
-            "output_tokens": response.usage.completion_tokens,
+            "input_tokens": input_tokens,
+            "output_tokens": output_tokens,
+            "cache_read_input_tokens": cached_tokens,
+            "reasoning_tokens": reasoning_tokens,
        }
    return {
        "input_tokens": 0,
        "output_tokens": 0,
+        "cache_read_input_tokens": 0,
+        "cache_creation_input_tokens": 0,
+        "reasoning_tokens": 0,
    }


@@ -75,24 +111,125 @@ def format_response_anthropic(response):

 def format_response_openai(response):
    output = []
-    for choice in response.choices:
-        if choice.message.content:
-            output.append(
-                {
-                    "content": choice.message.content,
-                    "role": choice.message.role,
-                }
-            )
+    if hasattr(response, "choices"):
+        for choice in response.choices:
+            # Handle Chat Completions response format
+            if hasattr(choice, "message") and choice.message and choice.message.content:
+                output.append(
+                    {
+                        "content": choice.message.content,
+                        "role": choice.message.role,
+                    }
+                )
+    # Handle Responses API format
+    if hasattr(response, "output"):
+        for item in response.output:
+            if item.type == "message":
+                # Extract text content from the content list
+                if hasattr(item, "content") and isinstance(item.content, list):
+                    for content_item in item.content:
+                        if (
+                            hasattr(content_item, "type")
+                            and content_item.type == "output_text"
+                            and hasattr(content_item, "text")
+                        ):
+                            output.append(
+                                {
+                                    "content": content_item.text,
+                                    "role": item.role,
+                                }
+                            )
+                        elif hasattr(content_item, "text"):
+                            output.append(
+                                {
+                                    "content": content_item.text,
+                                    "role": item.role,
+                                }
+                            )
+                        elif (
+                            hasattr(content_item, "type")
+                            and content_item.type == "input_image"
+                            and hasattr(content_item, "image_url")
+                        ):
+                            output.append(
+                                {
+                                    "content": {
+                                        "type": "image",
+                                        "image": content_item.image_url,
+                                    },
+                                    "role": item.role,
+                                }
+                            )
+                else:
+                    output.append(
+                        {
+                            "content": item.content,
+                            "role": item.role,
+                        }
+                    )
    return output


+def format_tool_calls(response, provider: str):
+    if provider == "anthropic":
+        if hasattr(response, "tools") and response.tools and len(response.tools) > 0:
+            return response.tools
+    elif provider == "openai":
+        # Handle both Chat Completions and Responses API
+        if hasattr(response, "choices") and response.choices:
+            # Check for tool_calls in message (Chat Completions format)
+            if (
+                hasattr(response.choices[0], "message")
+                and hasattr(response.choices[0].message, "tool_calls")
+                and response.choices[0].message.tool_calls
+            ):
+                return response.choices[0].message.tool_calls
+
+            # Check for tool_calls directly in response (Responses API format)
+            if hasattr(response.choices[0], "tool_calls") and response.choices[0].tool_calls:
+                return response.choices[0].tool_calls
+    return None
+
+
 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
+    messages: List[Dict[str, Any]] = []
+    if provider == "anthropic":
+        messages = kwargs.get("messages") or []
+        if kwargs.get("system") is None:
+            return messages
+        return [{"role": "system", "content": kwargs.get("system")}] + messages
+
+    # For OpenAI, handle both Chat Completions and Responses API
+    if kwargs.get("messages") is not None:
+        messages = list(kwargs.get("messages", []))
+
+    if kwargs.get("input") is not None:
+        input_data = kwargs.get("input")
+        if isinstance(input_data, list):
+            messages.extend(input_data)
+        else:
+            messages.append({"role": "user", "content": input_data})
+
+    # Check if system prompt is provided as a separate parameter
+    if kwargs.get("system") is not None:
+        has_system = any(msg.get("role") == "system" for msg in messages)
+        if not has_system:
+            messages = [{"role": "system", "content": kwargs.get("system")}] + messages
+
+    # For Responses API, add instructions to the system prompt if provided
+    if kwargs.get("instructions") is not None:
+        # Find the system message if it exists
+        system_idx = next((i for i, msg in enumerate(messages) if msg.get("role") == "system"), None)
+
+        if system_idx is not None:
+            # Append instructions to existing system message
+            system_content = messages[system_idx].get("content", "")
+            messages[system_idx]["content"] = f"{system_content}\n\n{kwargs.get('instructions')}"
+        else:
+            # Create a new system message with instructions
+            messages = [{"role": "system", "content": kwargs.get("instructions")}] + messages
+
+    return messages


 def call_llm_and_track_usage(
@@ -132,7 +269,7 @@ def call_llm_and_track_usage(
        latency = end_time - start_time

        if posthog_trace_id is None:
-            posthog_trace_id = uuid.uuid4()
+            posthog_trace_id = str(uuid.uuid4())

        if response and hasattr(response, "usage"):
            usage = get_usage(response, provider)
@@ -157,9 +294,28 @@ def call_llm_and_track_usage(
            **(error_params or {}),
        }

+        tool_calls = format_tool_calls(response, provider)
+        if tool_calls:
+            event_properties["$ai_tools"] = with_privacy_mode(ph_client, posthog_privacy_mode, tool_calls)
+
+        if usage.get("cache_read_input_tokens") is not None and usage.get("cache_read_input_tokens", 0) > 0:
+            event_properties["$ai_cache_read_input_tokens"] = usage.get("cache_read_input_tokens", 0)
+
+        if usage.get("cache_creation_input_tokens") is not None and usage.get("cache_creation_input_tokens", 0) > 0:
+            event_properties["$ai_cache_creation_input_tokens"] = usage.get("cache_creation_input_tokens", 0)
+
+        if usage.get("reasoning_tokens") is not None and usage.get("reasoning_tokens", 0) > 0:
+            event_properties["$ai_reasoning_tokens"] = usage.get("reasoning_tokens", 0)
+
        if posthog_distinct_id is None:
            event_properties["$process_person_profile"] = False

+        # Process instructions for Responses API
+        if provider == "openai" and kwargs.get("instructions") is not None:
+            event_properties["$ai_instructions"] = with_privacy_mode(
+                ph_client, posthog_privacy_mode, kwargs.get("instructions")
+            )
+
        # send the event to posthog
        if hasattr(ph_client, "capture") and callable(ph_client.capture):
            ph_client.capture(
@@ -208,7 +364,7 @@ async def call_llm_and_track_usage_async(
        latency = end_time - start_time

        if posthog_trace_id is None:
-            posthog_trace_id = uuid.uuid4()
+            posthog_trace_id = str(uuid.uuid4())

        if response and hasattr(response, "usage"):
            usage = get_usage(response, provider)
@@ -233,9 +389,25 @@ async def call_llm_and_track_usage_async(
            **(error_params or {}),
        }

+        tool_calls = format_tool_calls(response, provider)
+        if tool_calls:
+            event_properties["$ai_tools"] = with_privacy_mode(ph_client, posthog_privacy_mode, tool_calls)
+
+        if usage.get("cache_read_input_tokens") is not None and usage.get("cache_read_input_tokens", 0) > 0:
+            event_properties["$ai_cache_read_input_tokens"] = usage.get("cache_read_input_tokens", 0)
+
+        if usage.get("cache_creation_input_tokens") is not None and usage.get("cache_creation_input_tokens", 0) > 0:
+            event_properties["$ai_cache_creation_input_tokens"] = usage.get("cache_creation_input_tokens", 0)
+
        if posthog_distinct_id is None:
            event_properties["$process_person_profile"] = False

+        # Process instructions for Responses API
+        if provider == "openai" and kwargs.get("instructions") is not None:
+            event_properties["$ai_instructions"] = with_privacy_mode(
+                ph_client, posthog_privacy_mode, kwargs.get("instructions")
+            )
+
        # send the event to posthog
        if hasattr(ph_client, "capture") and callable(ph_client.capture):
            ph_client.capture(
@@ -2,11 +2,14 @@ import atexit
 import logging
 import numbers
 import os
+import platform
 import sys
 import warnings
 from datetime import datetime, timedelta
+from typing import Any
 from uuid import UUID, uuid4

+import distro  # For Linux OS detection
 from dateutil.tz import tzutc
 from six import string_types

@@ -29,6 +32,60 @@ ID_TYPES = (numbers.Number, string_types, UUID)
 MAX_DICT_SIZE = 50_000


+def get_os_info():
+    """
+    Returns standardized OS name and version information.
+    Similar to how user agent parsing works in JS.
+    """
+    os_name = ""
+    os_version = ""
+
+    platform_name = sys.platform
+
+    if platform_name.startswith("win"):
+        os_name = "Windows"
+        if hasattr(platform, "win32_ver"):
+            win_version = platform.win32_ver()[0]
+            if win_version:
+                os_version = win_version
+
+    elif platform_name == "darwin":
+        os_name = "Mac OS X"
+        if hasattr(platform, "mac_ver"):
+            mac_version = platform.mac_ver()[0]
+            if mac_version:
+                os_version = mac_version
+
+    elif platform_name.startswith("linux"):
+        os_name = "Linux"
+        linux_info = distro.info()
+        if linux_info["version"]:
+            os_version = linux_info["version"]
+
+    elif platform_name.startswith("freebsd"):
+        os_name = "FreeBSD"
+        if hasattr(platform, "release"):
+            os_version = platform.release()
+
+    else:
+        os_name = platform_name
+        if hasattr(platform, "release"):
+            os_version = platform.release()
+
+    return os_name, os_version
+
+
+def system_context() -> dict[str, Any]:
+    os_name, os_version = get_os_info()
+
+    return {
+        "$python_runtime": platform.python_implementation(),
+        "$python_version": "%s.%s.%s" % (sys.version_info[:3]),
+        "$os": os_name,
+        "$os_version": os_version,
+    }
+
+
 class Client(object):
    """Create a new PostHog client."""

@@ -231,7 +288,8 @@ class Client(object):
                stacklevel=2,
            )

-        properties = properties or {}
+        properties = {**(properties or {}), **system_context()}
+
        require("distinct_id", distinct_id, ID_TYPES)
        require("properties", properties, dict)
        require("event", event, string_types)
@@ -600,6 +658,21 @@ class Client(object):
                        "To use feature flags, please set a personal_api_key "
                        "More information: https://posthog.com/docs/api/overview",
                    )
+            elif e.status == 402:
+                self.log.warning(
+                    "[FEATURE FLAGS] PostHog feature flags quota limited, resetting feature flag data.  Learn more about billing limits at https://posthog.com/docs/billing/limits-alerts"
+                )
+                # Reset all feature flag data when quota limited
+                self.feature_flags = []
+                self.feature_flags_by_key = {}
+                self.group_type_mapping = {}
+                self.cohorts = {}
+
+                if self.debug:
+                    raise APIError(
+                        status=402,
+                        message="PostHog feature flags quota limited",
+                    )
            else:
                self.log.error(f"[FEATURE FLAGS] Error loading feature flags: {e}")
        except Exception as e:
@@ -822,7 +895,7 @@ class Client(object):
                    distinct_id, groups, person_properties, group_properties, disable_geoip
                )
                response = responses_and_payloads["featureFlags"].get(key, None)
-                payload = responses_and_payloads["featureFlagPayloads"].get(str(key).lower(), None)
+                payload = responses_and_payloads["featureFlagPayloads"].get(str(key), None)
            except Exception as e:
                self.log.exception(f"[FEATURE FLAGS] Unable to get feature flags and payloads: {e}")

@@ -875,10 +948,14 @@ class Client(object):
        if self.feature_flags_by_key is None:
            return payload

-        flag_definition = self.feature_flags_by_key.get(key) or {}
-        flag_filters = flag_definition.get("filters") or {}
-        flag_payloads = flag_filters.get("payloads") or {}
-        payload = flag_payloads.get(str(match_value).lower(), None)
+        flag_definition = self.feature_flags_by_key.get(key)
+        if flag_definition:
+            flag_filters = flag_definition.get("filters") or {}
+            flag_payloads = flag_filters.get("payloads") or {}
+            # For boolean flags, convert True to "true"
+            # For multivariate flags, use the variant string as-is
+            lookup_value = "true" if isinstance(match_value, bool) and match_value else str(match_value)
+            payload = flag_payloads.get(lookup_value, None)
        return payload

    def get_all_flags(
@@ -68,7 +68,21 @@ def _process_response(
    log = logging.getLogger("posthog")
    if res.status_code == 200:
        log.debug(success_message)
-        return res.json() if return_json else res
+        response = res.json() if return_json else res
+        # Handle quota limited decide responses by raising a specific error
+        # NB: other services also put entries into the quotaLimited key, but right now we only care about feature flags
+        # since most of the other services handle quota limiting in other places in the application.
+        if (
+            isinstance(response, dict)
+            and "quotaLimited" in response
+            and isinstance(response["quotaLimited"], list)
+            and "feature_flags" in response["quotaLimited"]
+        ):
+            log.warning(
+                "[FEATURE FLAGS] PostHog feature flags quota limited, resetting feature flag data.  Learn more about billing limits at https://posthog.com/docs/billing/limits-alerts"
+            )
+            raise QuotaLimitError(res.status_code, "Feature flags quota limited")
+        return response
    try:
        payload = res.json()
        log.debug("received response: %s", payload)
@@ -112,6 +126,10 @@ class APIError(Exception):
        return msg.format(self.message, self.status)


+class QuotaLimitError(APIError):
+    pass
+
+
 class DatetimeSerializer(json.JSONEncoder):
    def default(self, obj: Any):
        if isinstance(obj, (date, datetime)):
@@ -55,6 +55,28 @@ def mock_anthropic_stream():
    return stream_generator()


+@pytest.fixture
+def mock_anthropic_response_with_cached_tokens():
+    # Create a mock Usage object with cached_tokens in input_tokens_details
+    usage = Usage(
+        input_tokens=20,
+        output_tokens=10,
+        cache_read_input_tokens=15,
+        cache_creation_input_tokens=2,
+    )
+
+    return Message(
+        id="msg_123",
+        type="message",
+        role="assistant",
+        content=[{"type": "text", "text": "Test response"}],
+        model="claude-3-opus-20240229",
+        usage=usage,
+        stop_reason="end_turn",
+        stop_sequence=None,
+    )
+
+
 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)
@@ -339,3 +361,34 @@ def test_error(mock_client, mock_anthropic_response):
        props = call_args["properties"]
        assert props["$ai_is_error"] is True
        assert props["$ai_error"] == "Test error"
+
+
+def test_cached_tokens(mock_client, mock_anthropic_response_with_cached_tokens):
+    with patch("anthropic.resources.Messages.create", return_value=mock_anthropic_response_with_cached_tokens):
+        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_with_cached_tokens
+        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_cache_read_input_tokens"] == 15
+        assert props["$ai_cache_creation_input_tokens"] == 2
+        assert props["$ai_http_status"] == 200
+        assert props["foo"] == "bar"
+        assert isinstance(props["$ai_latency"], float)
@@ -1168,6 +1168,61 @@ async def test_async_anthropic_streaming(mock_client):
    assert isinstance(trace_props["$ai_output_state"], AIMessage)


+def test_metadata_tools(mock_client):
+    callbacks = CallbackHandler(mock_client)
+    run_id = uuid.uuid4()
+    tools = [
+        [
+            {
+                "type": "function",
+                "function": {
+                    "name": "foo",
+                    "description": "The foo.",
+                    "parameters": {
+                        "properties": {
+                            "bar": {
+                                "description": "The bar of foo.",
+                                "type": "string",
+                            },
+                        },
+                        "required": ["query_description", "query_kind"],
+                        "type": "object",
+                        "additionalProperties": False,
+                    },
+                    "strict": True,
+                },
+            }
+        ]
+    ]
+
+    with patch("time.time", return_value=1234567890):
+        callbacks._set_llm_metadata(
+            {"kwargs": {"openai_api_base": "https://us.posthog.com"}},
+            run_id,
+            messages=[{"role": "user", "content": "What's the weather like in SF?"}],
+            invocation_params={"temperature": 0.5, "tools": tools},
+            metadata={"ls_model_name": "hog-mini", "ls_provider": "posthog"},
+            name="test",
+        )
+    expected = GenerationMetadata(
+        model="hog-mini",
+        input=[{"role": "user", "content": "What's the weather like in SF?"}],
+        start_time=1234567890,
+        model_params={"temperature": 0.5},
+        provider="posthog",
+        base_url="https://us.posthog.com",
+        name="test",
+        tools=tools,
+        end_time=None,
+    )
+    assert callbacks._runs[run_id] == expected
+    with patch("time.time", return_value=1234567891):
+        run = callbacks._pop_run_metadata(run_id)
+    expected.end_time = 1234567891
+    assert run == expected
+    assert callbacks._runs == {}
+
+
 def test_tool_calls(mock_client):
    prompt = ChatPromptTemplate.from_messages([("user", "Foo")])
    model = FakeMessagesListChatModel(
@@ -1,12 +1,18 @@
+import json
 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.chat.chat_completion_chunk import ChatCompletionChunk
+from openai.types.chat.chat_completion_chunk import Choice as ChoiceChunk
+from openai.types.chat.chat_completion_chunk import ChoiceDelta, ChoiceDeltaToolCall, ChoiceDeltaToolCallFunction
+from openai.types.chat.chat_completion_message_tool_call import ChatCompletionMessageToolCall, Function
 from openai.types.completion_usage import CompletionUsage
 from openai.types.create_embedding_response import CreateEmbeddingResponse, Usage
 from openai.types.embedding import Embedding
+from openai.types.responses import Response, ResponseOutputMessage, ResponseOutputText, ResponseUsage

 from posthog.ai.openai import OpenAI

@@ -43,6 +49,48 @@ def mock_openai_response():
    )


+@pytest.fixture
+def mock_openai_response_with_responses_api():
+    return Response(
+        id="test",
+        model="gpt-4o-mini",
+        object="response",
+        created_at=1741476542,
+        status="completed",
+        error=None,
+        incomplete_details=None,
+        instructions=None,
+        max_output_tokens=None,
+        tools=[],
+        tool_choice="auto",
+        output=[
+            ResponseOutputMessage(
+                id="msg_123",
+                type="message",
+                role="assistant",
+                status="completed",
+                content=[
+                    ResponseOutputText(
+                        type="output_text",
+                        text="Test response",
+                        annotations=[],
+                    )
+                ],
+            )
+        ],
+        parallel_tool_calls=True,
+        previous_response_id=None,
+        usage=ResponseUsage(
+            input_tokens=10,
+            output_tokens=10,
+            output_tokens_details={"reasoning_tokens": 15},
+            total_tokens=20,
+        ),
+        user=None,
+        metadata={},
+    )
+
+
@pytest.fixture
 def mock_embedding_response():
    return CreateEmbeddingResponse(
@@ -62,6 +110,67 @@ def mock_embedding_response():
    )


+@pytest.fixture
+def mock_openai_response_with_cached_tokens():
+    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,
+            prompt_tokens_details={"cached_tokens": 15},
+        ),
+    )
+
+
+@pytest.fixture
+def mock_openai_response_with_tool_calls():
+    return ChatCompletion(
+        id="test",
+        model="gpt-4",
+        object="chat.completion",
+        created=int(time.time()),
+        choices=[
+            Choice(
+                finish_reason="tool_calls",
+                index=0,
+                message=ChatCompletionMessage(
+                    content="I'll check the weather for you.",
+                    role="assistant",
+                    tool_calls=[
+                        ChatCompletionMessageToolCall(
+                            id="call_abc123",
+                            type="function",
+                            function=Function(
+                                name="get_weather",
+                                arguments='{"location": "San Francisco", "unit": "celsius"}',
+                            ),
+                        )
+                    ],
+                ),
+            )
+        ],
+        usage=CompletionUsage(
+            completion_tokens=15,
+            prompt_tokens=20,
+            total_tokens=35,
+        ),
+    )
+
+
 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)
@@ -187,3 +296,279 @@ def test_error(mock_client, mock_openai_response):
        props = call_args["properties"]
        assert props["$ai_is_error"] is True
        assert props["$ai_error"] == "Test error"
+
+
+def test_cached_tokens(mock_client, mock_openai_response_with_cached_tokens):
+    with patch(
+        "openai.resources.chat.completions.Completions.create", return_value=mock_openai_response_with_cached_tokens
+    ):
+        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_with_cached_tokens
+        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_cache_read_input_tokens"] == 15
+        assert props["$ai_http_status"] == 200
+        assert props["foo"] == "bar"
+        assert isinstance(props["$ai_latency"], float)
+
+
+def test_tool_calls(mock_client, mock_openai_response_with_tool_calls):
+    with patch(
+        "openai.resources.chat.completions.Completions.create", return_value=mock_openai_response_with_tool_calls
+    ):
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+        response = client.chat.completions.create(
+            model="gpt-4",
+            messages=[{"role": "user", "content": "What's the weather in San Francisco?"}],
+            tools=[
+                {
+                    "type": "function",
+                    "function": {"name": "get_weather", "description": "Get weather", "parameters": {}},
+                }
+            ],
+            posthog_distinct_id="test-id",
+        )
+
+        assert response == mock_openai_response_with_tool_calls
+        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": "What's the weather in San Francisco?"}]
+        assert props["$ai_output_choices"] == [{"role": "assistant", "content": "I'll check the weather for you."}]
+
+        # Check that tool calls are properly captured
+        assert "$ai_tools" in props
+        tool_calls = props["$ai_tools"]
+        assert len(tool_calls) == 1
+
+        # Verify the tool call details
+        tool_call = tool_calls[0]
+        assert tool_call.id == "call_abc123"
+        assert tool_call.type == "function"
+        assert tool_call.function.name == "get_weather"
+
+        # Verify the arguments
+        arguments = tool_call.function.arguments
+        parsed_args = json.loads(arguments)
+        assert parsed_args == {"location": "San Francisco", "unit": "celsius"}
+
+        # Check token usage
+        assert props["$ai_input_tokens"] == 20
+        assert props["$ai_output_tokens"] == 15
+        assert props["$ai_http_status"] == 200
+
+
+def test_streaming_with_tool_calls(mock_client):
+    # Create mock tool call chunks that will be returned in sequence
+    tool_call_chunks = [
+        ChatCompletionChunk(
+            id="chunk1",
+            model="gpt-4",
+            object="chat.completion.chunk",
+            created=1234567890,
+            choices=[
+                ChoiceChunk(
+                    index=0,
+                    delta=ChoiceDelta(
+                        role="assistant",
+                        tool_calls=[
+                            ChoiceDeltaToolCall(
+                                index=0,
+                                id="call_abc123",
+                                type="function",
+                                function=ChoiceDeltaToolCallFunction(
+                                    name="get_weather",
+                                    arguments='{"location": "',
+                                ),
+                            )
+                        ],
+                    ),
+                    finish_reason=None,
+                )
+            ],
+        ),
+        ChatCompletionChunk(
+            id="chunk2",
+            model="gpt-4",
+            object="chat.completion.chunk",
+            created=1234567891,
+            choices=[
+                ChoiceChunk(
+                    index=0,
+                    delta=ChoiceDelta(
+                        tool_calls=[
+                            ChoiceDeltaToolCall(
+                                index=0,
+                                id="call_abc123",
+                                type="function",
+                                function=ChoiceDeltaToolCallFunction(
+                                    arguments='San Francisco"',
+                                ),
+                            )
+                        ],
+                    ),
+                    finish_reason=None,
+                )
+            ],
+        ),
+        ChatCompletionChunk(
+            id="chunk3",
+            model="gpt-4",
+            object="chat.completion.chunk",
+            created=1234567892,
+            choices=[
+                ChoiceChunk(
+                    index=0,
+                    delta=ChoiceDelta(
+                        tool_calls=[
+                            ChoiceDeltaToolCall(
+                                index=0,
+                                id="call_abc123",
+                                type="function",
+                                function=ChoiceDeltaToolCallFunction(
+                                    arguments=', "unit": "celsius"}',
+                                ),
+                            )
+                        ],
+                    ),
+                    finish_reason=None,
+                )
+            ],
+        ),
+        ChatCompletionChunk(
+            id="chunk4",
+            model="gpt-4",
+            object="chat.completion.chunk",
+            created=1234567893,
+            choices=[
+                ChoiceChunk(
+                    index=0,
+                    delta=ChoiceDelta(
+                        content="The weather in San Francisco is 15°C.",
+                    ),
+                    finish_reason=None,
+                )
+            ],
+            usage=CompletionUsage(
+                prompt_tokens=20,
+                completion_tokens=15,
+                total_tokens=35,
+            ),
+        ),
+    ]
+
+    # Mock the create method to return our chunks
+    with patch("openai.resources.chat.completions.Completions.create") as mock_create:
+        # Set up the mock to return our chunks when iterated
+        mock_create.return_value = tool_call_chunks
+
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+
+        # Call the streaming method
+        response_generator = client.chat.completions.create(
+            model="gpt-4",
+            messages=[{"role": "user", "content": "What's the weather in San Francisco?"}],
+            tools=[
+                {
+                    "type": "function",
+                    "function": {"name": "get_weather", "description": "Get weather", "parameters": {}},
+                }
+            ],
+            stream=True,
+            posthog_distinct_id="test-id",
+        )
+
+        # Consume the generator to trigger the event capture
+        chunks = list(response_generator)
+
+        # Verify the chunks were returned correctly
+        assert len(chunks) == 4
+        assert chunks == tool_call_chunks
+
+        # Verify the capture was called with the right arguments
+        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"
+
+        # Check that the tool calls were properly accumulated
+        assert "$ai_tools" in props
+        tool_calls = props["$ai_tools"]
+        assert len(tool_calls) == 1
+
+        # Verify the complete tool call was properly assembled
+        tool_call = tool_calls[0]
+        assert tool_call.id == "call_abc123"
+        assert tool_call.type == "function"
+        assert tool_call.function.name == "get_weather"
+
+        # Verify the arguments were concatenated correctly
+        arguments = tool_call.function.arguments
+        parsed_args = json.loads(arguments)
+        assert parsed_args == {"location": "San Francisco", "unit": "celsius"}
+
+        # Check that the content was also accumulated
+        assert props["$ai_output_choices"][0]["content"] == "The weather in San Francisco is 15°C."
+
+        # Check token usage
+        assert props["$ai_input_tokens"] == 20
+        assert props["$ai_output_tokens"] == 15
+
+
+# test responses api
+def test_responses_api(mock_client, mock_openai_response_with_responses_api):
+    with patch("openai.resources.responses.Responses.create", return_value=mock_openai_response_with_responses_api):
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+        response = client.responses.create(
+            model="gpt-4o-mini",
+            input="Hello",
+            posthog_distinct_id="test-id",
+            posthog_properties={"foo": "bar"},
+        )
+        assert response == mock_openai_response_with_responses_api
+        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-4o-mini"
+        assert props["$ai_input"] == [{"role": "user", "content": "Hello"}]
+        assert props["$ai_output_choices"] == [{"role": "assistant", "content": "Test response"}]
+        assert props["$ai_input_tokens"] == 10
+        assert props["$ai_output_tokens"] == 10
+        assert props["$ai_reasoning_tokens"] == 15
+        assert props["$ai_http_status"] == 200
+        assert props["foo"] == "bar"
+        assert isinstance(props["$ai_latency"], float)
@@ -5,8 +5,10 @@ from uuid import uuid4

 import mock
 import six
+from parameterized import parameterized

 from posthog.client import Client
+from posthog.request import APIError
 from posthog.test.test_utils import FAKE_TEST_API_KEY
 from posthog.version import VERSION

@@ -53,6 +55,11 @@ class TestClient(unittest.TestCase):
        self.assertEqual(msg["distinct_id"], "distinct_id")
        self.assertEqual(msg["properties"]["$lib"], "posthog-python")
        self.assertEqual(msg["properties"]["$lib_version"], VERSION)
+        # these will change between platforms so just asssert on presence here
+        assert msg["properties"]["$python_runtime"] == mock.ANY
+        assert msg["properties"]["$python_version"] == mock.ANY
+        assert msg["properties"]["$os"] == mock.ANY
+        assert msg["properties"]["$os_version"] == mock.ANY

    def test_basic_capture_with_uuid(self):
        client = self.client
@@ -100,7 +107,6 @@ class TestClient(unittest.TestCase):
        self.assertEqual(msg["properties"]["source"], "repo-name")

    def test_basic_capture_exception(self):
-
        with mock.patch.object(Client, "capture", return_value=None) as patch_capture:
            client = self.client
            exception = Exception("test exception")
@@ -128,7 +134,6 @@ class TestClient(unittest.TestCase):
            )

    def test_basic_capture_exception_with_distinct_id(self):
-
        with mock.patch.object(Client, "capture", return_value=None) as patch_capture:
            client = self.client
            exception = Exception("test exception")
@@ -156,7 +161,6 @@ class TestClient(unittest.TestCase):
            )

    def test_basic_capture_exception_with_correct_host_generation(self):
-
        with mock.patch.object(Client, "capture", return_value=None) as patch_capture:
            client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, host="https://aloha.com")
            exception = Exception("test exception")
@@ -184,7 +188,6 @@ class TestClient(unittest.TestCase):
            )

    def test_basic_capture_exception_with_correct_host_generation_for_server_hosts(self):
-
        with mock.patch.object(Client, "capture", return_value=None) as patch_capture:
            client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, host="https://app.posthog.com")
            exception = Exception("test exception")
@@ -212,7 +215,6 @@ class TestClient(unittest.TestCase):
            )

    def test_basic_capture_exception_with_no_exception_given(self):
-
        with mock.patch.object(Client, "capture", return_value=None) as patch_capture:
            client = self.client
            try:
@@ -249,10 +251,8 @@ class TestClient(unittest.TestCase):
            self.assertEqual(capture_call[2]["$exception_list"][0]["stacktrace"]["frames"][0]["in_app"], True)

    def test_basic_capture_exception_with_no_exception_happening(self):
-
        with mock.patch.object(Client, "capture", return_value=None) as patch_capture:
            with self.assertLogs("posthog", level="WARNING") as logs:
-
                client = self.client
                client.capture_exception()

@@ -384,6 +384,25 @@ class TestClient(unittest.TestCase):
        assert "$feature/false-flag" not in msg["properties"]
        assert "$active_feature_flags" not in msg["properties"]

+    @mock.patch("posthog.client.get")
+    def test_load_feature_flags_quota_limited(self, patch_get):
+        mock_response = {
+            "type": "quota_limited",
+            "detail": "You have exceeded your feature flag request quota",
+            "code": "payment_required",
+        }
+        patch_get.side_effect = APIError(402, mock_response["detail"])
+
+        client = Client(FAKE_TEST_API_KEY, personal_api_key="test")
+        with self.assertLogs("posthog", level="WARNING") as logs:
+            client._load_feature_flags()
+
+            self.assertEqual(client.feature_flags, [])
+            self.assertEqual(client.feature_flags_by_key, {})
+            self.assertEqual(client.group_type_mapping, {})
+            self.assertEqual(client.cohorts, {})
+            self.assertIn("PostHog feature flags quota limited", logs.output[0])
+
    @mock.patch("posthog.client.decide")
    def test_dont_override_capture_with_local_flags(self, patch_decide):
        patch_decide.return_value = {"featureFlags": {"beta-feature": "random-variant"}}
@@ -1104,3 +1123,92 @@ class TestClient(unittest.TestCase):
            group_properties={},
            disable_geoip=False,
        )
+
+    @parameterized.expand(
+        [
+            # name, sys_platform, version_info, expected_runtime, expected_version, expected_os, expected_os_version, platform_method, platform_return, distro_info
+            (
+                "macOS",
+                "darwin",
+                (3, 8, 10),
+                "MockPython",
+                "3.8.10",
+                "Mac OS X",
+                "10.15.7",
+                "mac_ver",
+                ("10.15.7", "", ""),
+                None,
+            ),
+            (
+                "Windows",
+                "win32",
+                (3, 8, 10),
+                "MockPython",
+                "3.8.10",
+                "Windows",
+                "10",
+                "win32_ver",
+                ("10", "", "", ""),
+                None,
+            ),
+            (
+                "Linux",
+                "linux",
+                (3, 8, 10),
+                "MockPython",
+                "3.8.10",
+                "Linux",
+                "20.04",
+                None,
+                None,
+                {"version": "20.04"},
+            ),
+        ]
+    )
+    def test_mock_system_context(
+        self,
+        _name,
+        sys_platform,
+        version_info,
+        expected_runtime,
+        expected_version,
+        expected_os,
+        expected_os_version,
+        platform_method,
+        platform_return,
+        distro_info,
+    ):
+        """Test that we can mock platform and sys for testing system_context"""
+        with mock.patch("posthog.client.platform") as mock_platform:
+            with mock.patch("posthog.client.sys") as mock_sys:
+                # Set up common mocks
+                mock_platform.python_implementation.return_value = expected_runtime
+                mock_sys.version_info = version_info
+                mock_sys.platform = sys_platform
+
+                # Set up platform-specific mocks
+                if platform_method:
+                    getattr(mock_platform, platform_method).return_value = platform_return
+
+                # Special handling for Linux which uses distro module
+                if sys_platform == "linux":
+                    # Directly patch the get_os_info function to return our expected values
+                    with mock.patch("posthog.client.get_os_info", return_value=(expected_os, expected_os_version)):
+                        from posthog.client import system_context
+
+                        context = system_context()
+                else:
+                    # Get system context for non-Linux platforms
+                    from posthog.client import system_context
+
+                    context = system_context()
+
+                # Verify results
+                expected_context = {
+                    "$python_runtime": expected_runtime,
+                    "$python_version": expected_version,
+                    "$os": expected_os,
+                    "$os_version": expected_os_version,
+                }
+
+                assert context == expected_context
@@ -4640,3 +4640,84 @@ class TestConsistency(unittest.TestCase):
                self.assertEqual(feature_flag_match, results[i])
            else:
                self.assertFalse(feature_flag_match)
+
+    @mock.patch("posthog.client.decide")
+    def test_feature_flag_case_sensitive(self, mock_decide):
+        mock_decide.return_value = {"featureFlags": {}}  # Ensure decide returns empty flags
+
+        client = Client(api_key=FAKE_TEST_API_KEY, personal_api_key=FAKE_TEST_API_KEY)
+        client.feature_flags = [
+            {
+                "id": 1,
+                "key": "Beta-Feature",
+                "active": True,
+                "filters": {
+                    "groups": [{"properties": [], "rollout_percentage": 100}],
+                },
+            }
+        ]
+
+        # Test that flag evaluation is case-sensitive
+        self.assertTrue(client.feature_enabled("Beta-Feature", "user1"))
+        self.assertFalse(client.feature_enabled("beta-feature", "user1"))
+        self.assertFalse(client.feature_enabled("BETA-FEATURE", "user1"))
+
+    @mock.patch("posthog.client.decide")
+    def test_feature_flag_payload_case_sensitive(self, mock_decide):
+        mock_decide.return_value = {
+            "featureFlags": {"Beta-Feature": True},
+            "featureFlagPayloads": {"Beta-Feature": {"some": "value"}},
+        }
+
+        client = Client(api_key=FAKE_TEST_API_KEY, personal_api_key=FAKE_TEST_API_KEY)
+        client.feature_flags = [
+            {
+                "id": 1,
+                "key": "Beta-Feature",
+                "active": True,
+                "filters": {
+                    "groups": [{"properties": [], "rollout_percentage": 100}],
+                    "payloads": {
+                        "true": {"some": "value"},
+                    },
+                },
+            }
+        ]
+
+        # Test that payload retrieval is case-sensitive
+        self.assertEqual(client.get_feature_flag_payload("Beta-Feature", "user1"), {"some": "value"})
+        self.assertIsNone(client.get_feature_flag_payload("beta-feature", "user1"))
+        self.assertIsNone(client.get_feature_flag_payload("BETA-FEATURE", "user1"))
+
+    @mock.patch("posthog.client.decide")
+    def test_feature_flag_case_sensitive_consistency(self, mock_decide):
+        mock_decide.return_value = {
+            "featureFlags": {"Beta-Feature": True},
+            "featureFlagPayloads": {"Beta-Feature": {"some": "value"}},
+        }
+
+        client = Client(api_key=FAKE_TEST_API_KEY, personal_api_key=FAKE_TEST_API_KEY)
+        client.feature_flags = [
+            {
+                "id": 1,
+                "key": "Beta-Feature",
+                "active": True,
+                "filters": {
+                    "groups": [{"properties": [], "rollout_percentage": 100}],
+                    "payloads": {
+                        "true": {"some": "value"},
+                    },
+                },
+            }
+        ]
+
+        # Test that flag evaluation and payload retrieval are consistently case-sensitive
+        # Only exact match should work
+        self.assertTrue(client.feature_enabled("Beta-Feature", "user1"))
+        self.assertEqual(client.get_feature_flag_payload("Beta-Feature", "user1"), {"some": "value"})
+
+        # Different cases should not match
+        test_cases = ["beta-feature", "BETA-FEATURE", "bEtA-FeAtUrE"]
+        for case in test_cases:
+            self.assertFalse(client.feature_enabled(case, "user1"))
+            self.assertIsNone(client.get_feature_flag_payload(case, "user1"))
@@ -2,10 +2,11 @@ import json
 import unittest
 from datetime import date, datetime

+import mock
 import pytest
 import requests

-from posthog.request import DatetimeSerializer, batch_post, determine_server_host
+from posthog.request import DatetimeSerializer, QuotaLimitError, batch_post, decide, determine_server_host
 from posthog.test.test_utils import TEST_API_KEY


@@ -44,6 +45,36 @@ class TestRequests(unittest.TestCase):
                "key", batch=[{"distinct_id": "distinct_id", "event": "python event", "type": "track"}], timeout=0.0001
            )

+    def test_quota_limited_response(self):
+        mock_response = requests.Response()
+        mock_response.status_code = 200
+        mock_response._content = json.dumps(
+            {
+                "quotaLimited": ["feature_flags"],
+                "featureFlags": {},
+                "featureFlagPayloads": {},
+                "errorsWhileComputingFlags": False,
+            }
+        ).encode("utf-8")
+
+        with mock.patch("posthog.request._session.post", return_value=mock_response):
+            with self.assertRaises(QuotaLimitError) as cm:
+                decide("fake_key", "fake_host")
+
+            self.assertEqual(cm.exception.status, 200)
+            self.assertEqual(cm.exception.message, "Feature flags quota limited")
+
+    def test_normal_decide_response(self):
+        mock_response = requests.Response()
+        mock_response.status_code = 200
+        mock_response._content = json.dumps(
+            {"featureFlags": {"flag1": True}, "featureFlagPayloads": {}, "errorsWhileComputingFlags": False}
+        ).encode("utf-8")
+
+        with mock.patch("posthog.request._session.post", return_value=mock_response):
+            response = decide("fake_key", "fake_host")
+            self.assertEqual(response["featureFlags"], {"flag1": True})
+

@pytest.mark.parametrize(
    "host, expected",
@@ -1,4 +1,4 @@
-VERSION = "3.14.1"
+VERSION = "3.20.0"

 if __name__ == "__main__":
    print(VERSION, end="")  # noqa: T201
@@ -20,6 +20,7 @@ install_requires = [
    "monotonic>=1.5",
    "backoff>=1.10.0",
    "python-dateutil>2.1",
+    "distro>=1.5.0",  # Required for Linux OS detection in Python 3.9+
 ]

 extras_require = {
@@ -57,6 +58,7 @@ extras_require = {
        "langchain-openai>=0.2.0",
        "langchain-anthropic>=0.2.0",
        "pydantic",
+        "parameterized>=0.8.1",
    ],
    "sentry": ["sentry-sdk", "django"],
    "langchain": ["langchain>=0.2.0"],
@@ -14,7 +14,14 @@ 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",
+    "distro>=1.5.0",  # Required for Linux OS detection in Python 3.9+
+]

 tests_require = ["mock>=2.0.0"]
Author	SHA1	Message	Date
Peter KirkhamandGitHub	5a4167d5ce	feat: add support for responses api (#205 ) * feat: add suppoort for responses api * fix: test * fix: black * fix: test - hopefully * fix: test - hopefully #2 * fix: test - hopefully #3 * fix: test - hopefully #4 * fix: greptaile catch * fix: mypy is not my friend * fix: isort usort weallsort * fix: noredef * fix: mypy baseline * fix: mypy * fix: mypy	2025-03-14 05:16:52 +00:00
David NewellandGitHub	332a6fffb6	fix: distro requirement for analytics package (#204 )	2025-03-12 14:12:15 +00:00
Peter KirkhamandGitHub	28a7d351ba	fix: azure open ai delta check (#203 )	2025-03-10 21:35:36 +00:00
Peter KirkhamandGitHub	8331af7a42	feat: cached tokens (#202 ) * feat: cached tokens * feat: add tool support * chore: local test * chore: isort black * chore: bump v * chore: remove import * fix: types * fix: black * fix: mypy unpacking of None * chore: mypy baseline * feat: mypy fix * fix: did things and stuff * fix: mypy yourpy whos py? * fix: things can be None * fix: move test * fix remove exampels from package * fix: losing my py	2025-03-06 22:37:21 +00:00
Dylan MartinandGitHub	f4c99714c3	chore(flags): improved some logs for quota limiting (#197 ) * haha okay * tests workin * format * use case-sensitive comparisons * omg LOL * fix tests * jeez * this will probably work * now do local eval * okay * yo * formatting * fix import order * type check * ai yi yi * code review * format * do it * merge conflict UGH * black formatting * bump version * correct changelog	2025-03-03 14:00:52 -05:00
Peter Kirkham GitHubgreptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>	7dc4cbb16b	feat: azure export w/ async (#200 ) * feat: azure export w/ async Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>	2025-02-28 20:40:09 +00:00
Michael MatlokaandGitHub	4cda646f03	feat(llm-observability): `$ai_tools` capture in Langchain (#199 )	2025-02-27 17:50:20 +00:00
Paul D'AmbraandGitHub	ea4e7fa16d	feat: add some platform info to events (#198 )	2025-02-26 12:26:17 +00:00
Peter KirkhamandGitHub	57a3e7470f	fix: async client (#196 )	2025-02-23 13:10:43 +00:00
Dylan MartinandGitHub	5e0f9e35c1	feat(feature-flags): support quota limiting for feature flags (#195 ) * haha okay * tests workin * format * use case-sensitive comparisons * omg LOL * fix tests * jeez * this will probably work * now do local eval * okay * yo * formatting * fix import order * type check * ai yi yi * code review * format	2025-02-21 15:45:51 -05:00
Dylan MartinandGitHub	337f7da7c5	fix(flags): remove `lower()` when evaluating feature flag payloads – these payloads are case-sensitive! (#191 ) * haha okay * tests workin * format * use case-sensitive comparisons * omg LOL * fix tests * jeez	2025-02-19 19:51:40 -05:00