update automatic retries to include read errors (#208)

in e.g. lambda environments the connection can time out between invocations,
this comes through to the client as a "RemoteDisconnected" error, which it
turns out urllib classifies as a "read" error not a connection error (as
it's possible to get this error after data has been sent)

.github/workflows/ci.yml

@@ -13,12 +13,12 @@ jobs:
               with:
                   fetch-depth: 1
 
-            - name: Set up Python 3.8
+            - name: Set up Python 3.11
               uses: actions/setup-python@v2
               with:
-                  python-version: 3.8
+                  python-version: 3.11.11
 
-            - uses: actions/cache@v1
+            - uses: actions/cache@v3
               with:
                   path: ~/.cache/pip
                   key: ${{ runner.os }}-pip-${{ hashFiles('setup.py') }}
@@ -33,28 +33,35 @@ jobs:
             - name: Check formatting with black
               run: |
                   black --check .
-            
+
             - name: Lint with flake8
               run: |
-                  flake8 posthog --ignore E501
+                  flake8 posthog --ignore E501,W503
 
             - name: Check import order with isort
               run: |
                   isort --check-only .
 
+            - name: Check types with mypy
+              run: |
+                  mypy --no-site-packages --config-file mypy.ini  . | mypy-baseline filter
+
     tests:
-        name: Python tests
+        name: Python ${{ matrix.python-version }} tests
         runs-on: ubuntu-latest
+        strategy:
+            matrix:
+                python-version: ['3.9', '3.10', '3.11', '3.12', '3.13']
 
         steps:
-            - uses: actions/checkout@v1
+            - uses: actions/checkout@v2
               with:
                   fetch-depth: 1
 
-            - name: Set up Python 3.7
-              uses: actions/setup-python@v1
+            - name: Set up Python ${{ matrix.python-version }}
+              uses: actions/setup-python@v2
               with:
-                  python-version: 3.7
+                  python-version: ${{ matrix.python-version }}
 
             - name: Install requirements.txt dependencies with pip
               run: |

.gitignore

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

CHANGELOG.md

@@ -1,3 +1,128 @@
+## 3.23.0 – 2025-03-26
+
+1. Expand automatic retries to include read errors (e.g. RemoteDisconnected)
+
+## 3.22.0 – 2025-03-26
+
+1. Add more information to `$feature_flag_called` events.
+2. Support for the `/decide?v=3` endpoint which contains more information about feature flags.
+
+## 3.21.0 – 2025-03-17
+
+1. Support serializing dataclasses.
+
+## 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
+
+## 3.13.0 - 2025-02-12
+
+1. Automatically retry connection errors
+
+## 3.12.1 - 2025-02-11
+
+1. Fix mypy support for 3.12.0
+2. Deprecate `is_simple_flag`
+
+## 3.12.0 - 2025-02-11
+
+1. Add support for OpenAI beta parse API.
+2. Deprecate `context` parameter
+
+## 3.11.1 - 2025-02-06
+
+1. Fix LangChain callback handler to capture parent run ID.
+
+## 3.11.0 - 2025-01-28
+
+1. Add the `$ai_span` event to the LangChain callback handler to capture the input and output of intermediary chains.
+
+   > LLM observability naming change: event property `$ai_trace_name` is now `$ai_span_name`.
+
+2. Fix serialiazation of Pydantic models in methods.
+
+## 3.10.0 - 2025-01-24
+
+1. Add `$ai_error` and `$ai_is_error` properties to LangChain callback handler, OpenAI, and Anthropic.
+
+## 3.9.3 - 2025-01-23
+
+1. Fix capturing of multiple traces in the LangChain callback handler.
+
+## 3.9.2 - 2025-01-22
+
+1. Fix importing of LangChain callback handler under certain circumstances.
+
+## 3.9.0 - 2025-01-22
+
+1. Add `$ai_trace` event emission to LangChain callback handler.
+
+## 3.8.4 - 2025-01-17
+
+1. Add Anthropic support for LLM Observability.
+2. Update LLM Observability to use output_choices.
+
+## 3.8.3 - 2025-01-14
+
+1. Fix setuptools to include the `posthog.ai.openai` and `posthog.ai.langchain` packages for the `posthoganalytics` package.
+
+## 3.8.2 - 2025-01-14
+
+1. Fix setuptools to include the `posthog.ai.openai` and `posthog.ai.langchain` packages.
+
+## 3.8.1 - 2025-01-14
+
+1. Add LLM Observability with support for OpenAI and Langchain callbacks.
+
+## 3.7.5 - 2025-01-03
+
+1. Add `distinct_id` to group_identify
+
 ## 3.7.4 - 2024-11-25
 
 1. Fix bug where this SDK incorrectly sent feature flag events with null values when calling `get_feature_flag_payload`.
@@ -78,6 +203,7 @@
 
 1. Update type hints for module variables to work with newer versions of mypy
 
+
 ## 3.3.3 - 2024-01-26
 
 1. Remove new relative date operators, combine into regular date operators

CODEOWNERS

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

Makefile

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

README.md

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

example.py

@@ -1,10 +1,15 @@
 # PostHog Python library example
-
-# Import the library
-import time
+import argparse
 
 import posthog
 
+# Add argument parsing
+parser = argparse.ArgumentParser(description="PostHog Python library example")
+parser.add_argument(
+    "--flag", default="person-on-events-enabled", help="Feature flag key to check (default: person-on-events-enabled)"
+)
+args = parser.parse_args()
+
 posthog.debug = True
 
 # You can find this key on the /setup page in PostHog
@@ -18,7 +23,7 @@ posthog.poll_interval = 10
 
 print(
     posthog.feature_enabled(
-        "person-on-events-enabled",
+        args.flag,  # Use the flag from command line arguments
         "12345",
         groups={"organization": str("0182ee91-8ef7-0000-4cb9-fedc5f00926a")},
         group_properties={
@@ -96,6 +101,7 @@ print(
         "distinct_id_random_22", person_properties={"$geoip_city_name": "Sydney"}, only_evaluate_locally=True
     )
 )
+print(posthog.get_remote_config_payload("encrypted_payload_flag_key"))
 
 
 posthog.shutdown()

mypy-baseline.txt

@@ -0,0 +1,41 @@
+posthog/utils.py:0: error: Library stubs not installed for "six"  [import-untyped]
+posthog/utils.py:0: error: Library stubs not installed for "dateutil.tz"  [import-untyped]
+posthog/utils.py:0: error: Statement is unreachable  [unreachable]
+posthog/utils.py:0: error: Argument 1 to "join" of "str" has incompatible type "AttributeError"; expected "Iterable[str]"  [arg-type]
+posthog/request.py:0: error: Library stubs not installed for "requests"  [import-untyped]
+posthog/request.py:0: note: Hint: "python3 -m pip install types-requests"
+posthog/request.py:0: error: Library stubs not installed for "dateutil.tz"  [import-untyped]
+posthog/request.py:0: error: Incompatible types in assignment (expression has type "bytes", variable has type "str")  [assignment]
+posthog/consumer.py:0: error: Name "Empty" already defined (possibly by an import)  [no-redef]
+posthog/consumer.py:0: error: Need type annotation for "items" (hint: "items: list[<type>] = ...")  [var-annotated]
+posthog/consumer.py:0: error: Unsupported operand types for <= ("int" and "str")  [operator]
+posthog/consumer.py:0: note: Right operand is of type "int | str"
+posthog/consumer.py:0: error: Unsupported operand types for < ("str" and "int")  [operator]
+posthog/consumer.py:0: note: Left operand is of type "int | str"
+posthog/feature_flags.py:0: error: Library stubs not installed for "dateutil"  [import-untyped]
+posthog/feature_flags.py:0: error: Library stubs not installed for "dateutil.relativedelta"  [import-untyped]
+posthog/feature_flags.py:0: error: Unused "type: ignore" comment  [unused-ignore]
+posthog/client.py:0: error: Library stubs not installed for "dateutil.tz"  [import-untyped]
+posthog/client.py:0: note: Hint: "python3 -m pip install types-python-dateutil"
+posthog/client.py:0: note: (or run "mypy --install-types" to install all missing stub packages)
+posthog/client.py:0: note: See https://mypy.readthedocs.io/en/stable/running_mypy.html#missing-imports
+posthog/client.py:0: error: Library stubs not installed for "six"  [import-untyped]
+posthog/client.py:0: note: Hint: "python3 -m pip install types-six"
+posthog/client.py:0: error: Name "queue" already defined (by an import)  [no-redef]
+posthog/client.py:0: error: Need type annotation for "queue"  [var-annotated]
+posthog/client.py:0: error: Item "None" of "Any | None" has no attribute "get"  [union-attr]
+simulator.py:0: error: Unexpected keyword argument "anonymous_id" for "capture"  [call-arg]
+posthog/__init__.py:0: note: "capture" defined here
+simulator.py:0: error: Unexpected keyword argument "anonymous_id" for "identify"  [call-arg]
+posthog/__init__.py:0: note: "identify" defined here
+simulator.py:0: error: Unexpected keyword argument "traits" for "identify"  [call-arg]
+posthog/__init__.py:0: note: "identify" defined here
+example.py:0: error: Statement is unreachable  [unreachable]
+posthog/sentry/posthog_integration.py:0: error: Statement is unreachable  [unreachable]
+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: 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"?
+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]

mypy.ini

@@ -0,0 +1,38 @@
+[mypy]
+python_version = 3.11
+plugins =
+    pydantic.mypy
+strict_optional = True
+no_implicit_optional = True
+warn_unused_ignores = True
+check_untyped_defs = True
+warn_unreachable = True
+strict_equality = True
+ignore_missing_imports = True
+
+[mypy-django.*]
+ignore_missing_imports = True
+
+[mypy-sentry_sdk.*]
+ignore_missing_imports = True
+
+[mypy-posthog.test.*]
+ignore_errors = True
+
+[mypy-posthog.*.test.*]
+ignore_errors = True
+
+[mypy-openai.*]
+ignore_missing_imports = True
+
+[mypy-langchain.*]
+ignore_missing_imports = True
+
+[mypy-langchain_core.*]
+ignore_missing_imports = True
+
+[mypy-anthropic.*]
+ignore_missing_imports = True
+
+[mypy-httpx.*]
+ignore_missing_imports = True

posthog/__init__.py

@@ -1,8 +1,10 @@
 import datetime  # noqa: F401
+import warnings
 from typing import Callable, Dict, List, Optional, Tuple  # noqa: F401
 
 from posthog.client import Client
 from posthog.exception_capture import Integrations  # noqa: F401
+from posthog.types import FeatureFlag, FlagsAndPayloads
 from posthog.version import VERSION
 
 __version__ = VERSION
@@ -26,6 +28,8 @@ enable_exception_autocapture = False  # type: bool
 exception_autocapture_integrations = []  # type: List[Integrations]
 # Used to determine in app paths for exception autocapture. Defaults to the current working directory
 project_root = None  # type: Optional[str]
+# Used for our AI observability feature to not capture any prompt or output just usage + metadata
+privacy_mode = False  # type: bool
 
 default_client = None  # type: Optional[Client]
 
@@ -62,6 +66,14 @@ def capture(
     posthog.capture('distinct id', 'purchase', groups={'company': 'id:5'})
     ```
     """
+
+    if context is not None:
+        warnings.warn(
+            "The 'context' parameter is deprecated and will be removed in a future version.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
     return _proxy(
         "capture",
         distinct_id=distinct_id,
@@ -100,6 +112,14 @@ def identify(
     })
     ```
     """
+
+    if context is not None:
+        warnings.warn(
+            "The 'context' parameter is deprecated and will be removed in a future version.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
     return _proxy(
         "identify",
         distinct_id=distinct_id,
@@ -135,6 +155,14 @@ def set(
      })
      ```
     """
+
+    if context is not None:
+        warnings.warn(
+            "The 'context' parameter is deprecated and will be removed in a future version.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
     return _proxy(
         "set",
         distinct_id=distinct_id,
@@ -170,6 +198,14 @@ def set_once(
      })
      ```
     """
+
+    if context is not None:
+        warnings.warn(
+            "The 'context' parameter is deprecated and will be removed in a future version.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
     return _proxy(
         "set_once",
         distinct_id=distinct_id,
@@ -206,6 +242,14 @@ def group_identify(
      })
      ```
     """
+
+    if context is not None:
+        warnings.warn(
+            "The 'context' parameter is deprecated and will be removed in a future version.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
     return _proxy(
         "group_identify",
         group_type=group_type,
@@ -243,6 +287,14 @@ def alias(
     posthog.alias('anonymous session id', 'distinct id')
     ```
     """
+
+    if context is not None:
+        warnings.warn(
+            "The 'context' parameter is deprecated and will be removed in a future version.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
     return _proxy(
         "alias",
         previous_id=previous_id,
@@ -286,6 +338,14 @@ def capture_exception(
 
     ```
     """
+
+    if context is not None:
+        warnings.warn(
+            "The 'context' parameter is deprecated and will be removed in a future version.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
     return _proxy(
         "capture_exception",
         exception=exception,
@@ -344,7 +404,7 @@ def get_feature_flag(
     only_evaluate_locally=False,  # type: bool
     send_feature_flag_events=True,  # type: bool
     disable_geoip=None,  # type: Optional[bool]
-):
+) -> Optional[FeatureFlag]:
     """
     Get feature flag variant for users. Used with experiments.
     Example:
@@ -387,7 +447,7 @@ def get_all_flags(
     group_properties={},  # type: dict
     only_evaluate_locally=False,  # type: bool
     disable_geoip=None,  # type: Optional[bool]
-):
+) -> Optional[dict[str, FeatureFlag]]:
     """
     Get all flags for a given user.
     Example:
@@ -418,7 +478,7 @@ def get_feature_flag_payload(
     only_evaluate_locally=False,
     send_feature_flag_events=True,
     disable_geoip=None,  # type: Optional[bool]
-):
+) -> Optional[str]:
     return _proxy(
         "get_feature_flag_payload",
         key=key,
@@ -433,6 +493,26 @@ def get_feature_flag_payload(
     )
 
 
+def get_remote_config_payload(
+    key,  # type: str
+):
+    """Get the payload for a remote config feature flag.
+
+    Args:
+        key: The key of the feature flag
+
+    Returns:
+        The payload associated with the feature flag. If payload is encrypted, the return value will decrypted
+
+    Note:
+        Requires personal_api_key to be set for authentication
+    """
+    return _proxy(
+        "get_remote_config_payload",
+        key=key,
+    )
+
+
 def get_all_flags_and_payloads(
     distinct_id,
     groups={},
@@ -440,7 +520,7 @@ def get_all_flags_and_payloads(
     group_properties={},
     only_evaluate_locally=False,
     disable_geoip=None,  # type: Optional[bool]
-):
+) -> FlagsAndPayloads:
     return _proxy(
         "get_all_flags_and_payloads",
         distinct_id=distinct_id,

posthog/ai/anthropic/__init__.py

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

posthog/ai/anthropic/anthropic.py

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

posthog/ai/anthropic/anthropic_async.py

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

posthog/ai/anthropic/anthropic_providers.py

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

posthog/ai/langchain/__init__.py

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

posthog/ai/langchain/callbacks.py

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

posthog/ai/openai/__init__.py

@@ -0,0 +1,5 @@
+from .openai import OpenAI
+from .openai_async import AsyncOpenAI
+from .openai_providers import AsyncAzureOpenAI, AzureOpenAI
+
+__all__ = ["OpenAI", "AsyncOpenAI", "AzureOpenAI", "AsyncAzureOpenAI"]

posthog/ai/openai/openai.py

@@ -0,0 +1,494 @@
+import time
+import uuid
+from typing import Any, Dict, List, Optional
+
+try:
+    import openai
+    import openai.resources
+except ImportError:
+    raise ModuleNotFoundError("Please install the OpenAI SDK to use this feature: 'pip install openai'")
+
+from posthog.ai.utils import call_llm_and_track_usage, get_model_params, with_privacy_mode
+from posthog.client import Client as PostHogClient
+
+
+class OpenAI(openai.OpenAI):
+    """
+    A wrapper around the OpenAI SDK that automatically sends LLM usage events to PostHog.
+    """
+
+    _ph_client: PostHogClient
+
+    def __init__(self, posthog_client: PostHogClient, **kwargs):
+        """
+        Args:
+            api_key: OpenAI API key.
+            posthog_client: If provided, events will be captured via this client instead
+                            of the global posthog.
+            **openai_config: Any additional keyword args to set on openai (e.g. organization="xxx").
+        """
+        super().__init__(**kwargs)
+        self._ph_client = posthog_client
+        self.chat = WrappedChat(self)
+        self.embeddings = WrappedEmbeddings(self)
+        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):
+    _client: OpenAI
+
+    @property
+    def completions(self):
+        return WrappedCompletions(self._client)
+
+
+class WrappedCompletions(openai.resources.chat.completions.Completions):
+    _client: OpenAI
+
+    def create(
+        self,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: bool = False,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        if posthog_trace_id is None:
+            posthog_trace_id = 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] = {}
+        accumulated_content = []
+        accumulated_tools = {}
+        if "stream_options" not in kwargs:
+            kwargs["stream_options"] = {}
+        kwargs["stream_options"]["include_usage"] = True
+        response = super().create(**kwargs)
+
+        def generator():
+            nonlocal usage_stats
+            nonlocal accumulated_content
+            nonlocal accumulated_tools
+
+            try:
+                for chunk in response:
+                    if hasattr(chunk, "usage") and chunk.usage:
+                        usage_stats = {
+                            k: getattr(chunk.usage, k, 0)
+                            for k in [
+                                "prompt_tokens",
+                                "completion_tokens",
+                                "total_tokens",
+                            ]
+                        }
+
+                        # 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:
+                        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
+                self._capture_streaming_event(
+                    posthog_distinct_id,
+                    posthog_trace_id,
+                    posthog_properties,
+                    posthog_privacy_mode,
+                    posthog_groups,
+                    kwargs,
+                    usage_stats,
+                    latency,
+                    output,
+                    tools,
+                )
+
+        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("messages")),
+            "$ai_output_choices": with_privacy_mode(
+                self._client._ph_client,
+                posthog_privacy_mode,
+                [{"content": output, "role": "assistant"}],
+            ),
+            "$ai_http_status": 200,
+            "$ai_input_tokens": usage_stats.get("prompt_tokens", 0),
+            "$ai_output_tokens": usage_stats.get("completion_tokens", 0),
+            "$ai_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 WrappedEmbeddings(openai.resources.embeddings.Embeddings):
+    _client: OpenAI
+
+    def create(
+        self,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: bool = False,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        """
+        Create an embedding using OpenAI's 'embeddings.create' method, but also track usage in PostHog.
+
+        Args:
+            posthog_distinct_id: Optional ID to associate with the usage event.
+            posthog_trace_id: Optional trace UUID for linking events.
+            posthog_properties: Optional dictionary of extra properties to include in the event.
+            **kwargs: Any additional parameters for the OpenAI Embeddings API.
+
+        Returns:
+            The response from OpenAI's embeddings.create call.
+        """
+        if posthog_trace_id is None:
+            posthog_trace_id = str(uuid.uuid4())
+
+        start_time = time.time()
+        response = super().create(**kwargs)
+        end_time = time.time()
+
+        # Extract usage statistics if available
+        usage_stats = {}
+        if hasattr(response, "usage") and response.usage:
+            usage_stats = {
+                "prompt_tokens": getattr(response.usage, "prompt_tokens", 0),
+                "total_tokens": getattr(response.usage, "total_tokens", 0),
+            }
+
+        latency = end_time - start_time
+
+        # Build the event properties
+        event_properties = {
+            "$ai_provider": "openai",
+            "$ai_model": kwargs.get("model"),
+            "$ai_input": with_privacy_mode(self._client._ph_client, posthog_privacy_mode, kwargs.get("input")),
+            "$ai_http_status": 200,
+            "$ai_input_tokens": usage_stats.get("prompt_tokens", 0),
+            "$ai_latency": latency,
+            "$ai_trace_id": posthog_trace_id,
+            "$ai_base_url": str(self._client.base_url),
+            **(posthog_properties or {}),
+        }
+
+        if posthog_distinct_id is None:
+            event_properties["$process_person_profile"] = False
+
+        # Send capture event for embeddings
+        if hasattr(self._client._ph_client, "capture"):
+            self._client._ph_client.capture(
+                distinct_id=posthog_distinct_id or posthog_trace_id,
+                event="$ai_embedding",
+                properties=event_properties,
+                groups=posthog_groups,
+            )
+
+        return response
+
+
+class WrappedBeta(openai.resources.beta.Beta):
+    _client: OpenAI
+
+    @property
+    def chat(self):
+        return WrappedBetaChat(self._client)
+
+
+class WrappedBetaChat(openai.resources.beta.chat.Chat):
+    _client: OpenAI
+
+    @property
+    def completions(self):
+        return WrappedBetaCompletions(self._client)
+
+
+class WrappedBetaCompletions(openai.resources.beta.chat.completions.Completions):
+    _client: OpenAI
+
+    def parse(
+        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,
+    ):
+        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().parse,
+            **kwargs,
+        )

posthog/ai/openai/openai_async.py

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

posthog/ai/openai/openai_providers.py

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

posthog/ai/utils.py

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

posthog/client.py

@@ -2,10 +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, Optional, Union
 from uuid import UUID, uuid4
 
+import distro  # For Linux OS detection
 from dateutil.tz import tzutc
 from six import string_types
 
@@ -14,7 +18,18 @@ from posthog.exception_capture import ExceptionCapture
 from posthog.exception_utils import exc_info_from_error, exceptions_from_error_tuple, handle_in_app
 from posthog.feature_flags import InconclusiveMatchError, match_feature_flag_properties
 from posthog.poller import Poller
-from posthog.request import DEFAULT_HOST, APIError, batch_post, decide, determine_server_host, get
+from posthog.request import DEFAULT_HOST, APIError, batch_post, decide, determine_server_host, get, remote_config
+from posthog.types import (
+    DecideResponse,
+    FeatureFlag,
+    FlagMetadata,
+    FlagsAndPayloads,
+    FlagValue,
+    normalize_decide_response,
+    to_flags_and_payloads,
+    to_payloads,
+    to_values,
+)
 from posthog.utils import SizeLimitedDict, clean, guess_timezone, remove_trailing_slash
 from posthog.version import VERSION
 
@@ -28,6 +43,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."""
 
@@ -59,6 +128,7 @@ class Client(object):
         enable_exception_autocapture=False,
         exception_autocapture_integrations=None,
         project_root=None,
+        privacy_mode=False,
     ):
         self.queue = queue.Queue(max_queue_size)
 
@@ -76,7 +146,7 @@ class Client(object):
         self.host = determine_server_host(host)
         self.gzip = gzip
         self.timeout = timeout
-        self.feature_flags = None
+        self._feature_flags = None  # private variable to store flags
         self.feature_flags_by_key = None
         self.group_type_mapping = None
         self.cohorts = None
@@ -91,6 +161,7 @@ class Client(object):
         self.enable_exception_autocapture = enable_exception_autocapture
         self.exception_autocapture_integrations = exception_autocapture_integrations
         self.exception_capture = None
+        self.privacy_mode = privacy_mode
 
         if project_root is None:
             try:
@@ -144,15 +215,38 @@ class Client(object):
                 if send:
                     consumer.start()
 
+    @property
+    def feature_flags(self):
+        """
+        Get the local evaluation feature flags.
+        """
+        return self._feature_flags
+
+    @feature_flags.setter
+    def feature_flags(self, flags):
+        """
+        Set the local evaluation feature flags.
+        """
+        self._feature_flags = flags or []
+        self.feature_flags_by_key = {flag["key"]: flag for flag in self._feature_flags if flag.get("key") is not None}
+        assert (
+            self.feature_flags_by_key is not None
+        ), "feature_flags_by_key should be initialized when feature_flags is set"
+
     def identify(self, distinct_id=None, properties=None, context=None, timestamp=None, uuid=None, disable_geoip=None):
+        if context is not None:
+            warnings.warn(
+                "The 'context' parameter is deprecated and will be removed in a future version.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
         properties = properties or {}
-        context = context or {}
         require("distinct_id", distinct_id, ID_TYPES)
         require("properties", properties, dict)
 
         msg = {
             "timestamp": timestamp,
-            "context": context,
             "distinct_id": distinct_id,
             "$set": properties,
             "event": "$identify",
@@ -163,26 +257,34 @@ class Client(object):
 
     def get_feature_variants(
         self, distinct_id, groups=None, person_properties=None, group_properties=None, disable_geoip=None
-    ):
+    ) -> dict[str, Union[bool, str]]:
+        """
+        Get feature flag variants for a distinct_id by calling decide.
+        """
         resp_data = self.get_decide(distinct_id, groups, person_properties, group_properties, disable_geoip)
-        return resp_data["featureFlags"]
+        return to_values(resp_data) or {}
 
     def get_feature_payloads(
         self, distinct_id, groups=None, person_properties=None, group_properties=None, disable_geoip=None
-    ):
+    ) -> dict[str, str]:
+        """
+        Get feature flag payloads for a distinct_id by calling decide.
+        """
         resp_data = self.get_decide(distinct_id, groups, person_properties, group_properties, disable_geoip)
-        return resp_data["featureFlagPayloads"]
+        return to_payloads(resp_data) or {}
 
     def get_feature_flags_and_payloads(
         self, distinct_id, groups=None, person_properties=None, group_properties=None, disable_geoip=None
-    ):
-        resp_data = self.get_decide(distinct_id, groups, person_properties, group_properties, disable_geoip)
-        return {
-            "featureFlags": resp_data["featureFlags"],
-            "featureFlagPayloads": resp_data["featureFlagPayloads"],
-        }
+    ) -> FlagsAndPayloads:
+        """
+        Get feature flags and payloads for a distinct_id by calling decide.
+        """
+        resp = self.get_decide(distinct_id, groups, person_properties, group_properties, disable_geoip)
+        return to_flags_and_payloads(resp)
 
-    def get_decide(self, distinct_id, groups=None, person_properties=None, group_properties=None, disable_geoip=None):
+    def get_decide(
+        self, distinct_id, groups=None, person_properties=None, group_properties=None, disable_geoip=None
+    ) -> DecideResponse:
         require("distinct_id", distinct_id, ID_TYPES)
 
         if disable_geoip is None:
@@ -202,7 +304,7 @@ class Client(object):
         }
         resp_data = decide(self.api_key, self.host, timeout=self.feature_flags_request_timeout_seconds, **request_data)
 
-        return resp_data
+        return normalize_decide_response(resp_data)
 
     def capture(
         self,
@@ -216,8 +318,15 @@ class Client(object):
         send_feature_flags=False,
         disable_geoip=None,
     ):
-        properties = properties or {}
-        context = context or {}
+        if context is not None:
+            warnings.warn(
+                "The 'context' parameter is deprecated and will be removed in a future version.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
+        properties = {**(properties or {}), **system_context()}
+
         require("distinct_id", distinct_id, ID_TYPES)
         require("properties", properties, dict)
         require("event", event, string_types)
@@ -225,7 +334,6 @@ class Client(object):
         msg = {
             "properties": properties,
             "timestamp": timestamp,
-            "context": context,
             "distinct_id": distinct_id,
             "event": event,
             "uuid": uuid,
@@ -235,24 +343,24 @@ class Client(object):
             require("groups", groups, dict)
             msg["properties"]["$groups"] = groups
 
-        extra_properties = {}
-        feature_variants = {}
+        extra_properties: dict[str, Any] = {}
+        feature_variants: Optional[dict[str, Union[bool, str]]] = {}
         if send_feature_flags:
             try:
                 feature_variants = self.get_feature_variants(distinct_id, groups, disable_geoip=disable_geoip)
             except Exception as e:
                 self.log.exception(f"[FEATURE FLAGS] Unable to get feature variants: {e}")
 
-        elif self.feature_flags:
+        elif self.feature_flags and event != "$feature_flag_called":
             # Local evaluation is enabled, flags are loaded, so try and get all flags we can without going to the server
             feature_variants = self.get_all_flags(
                 distinct_id, groups=(groups or {}), disable_geoip=disable_geoip, only_evaluate_locally=True
             )
 
-        for feature, variant in feature_variants.items():
+        for feature, variant in (feature_variants or {}).items():
             extra_properties[f"$feature/{feature}"] = variant
 
-        active_feature_flags = [key for (key, value) in feature_variants.items() if value is not False]
+        active_feature_flags = [key for (key, value) in (feature_variants or {}).items() if value is not False]
         if active_feature_flags:
             extra_properties["$active_feature_flags"] = active_feature_flags
 
@@ -262,14 +370,19 @@ class Client(object):
         return self._enqueue(msg, disable_geoip)
 
     def set(self, distinct_id=None, properties=None, context=None, timestamp=None, uuid=None, disable_geoip=None):
+        if context is not None:
+            warnings.warn(
+                "The 'context' parameter is deprecated and will be removed in a future version.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
         properties = properties or {}
-        context = context or {}
         require("distinct_id", distinct_id, ID_TYPES)
         require("properties", properties, dict)
 
         msg = {
             "timestamp": timestamp,
-            "context": context,
             "distinct_id": distinct_id,
             "$set": properties,
             "event": "$set",
@@ -279,14 +392,19 @@ class Client(object):
         return self._enqueue(msg, disable_geoip)
 
     def set_once(self, distinct_id=None, properties=None, context=None, timestamp=None, uuid=None, disable_geoip=None):
+        if context is not None:
+            warnings.warn(
+                "The 'context' parameter is deprecated and will be removed in a future version.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
         properties = properties or {}
-        context = context or {}
         require("distinct_id", distinct_id, ID_TYPES)
         require("properties", properties, dict)
 
         msg = {
             "timestamp": timestamp,
-            "context": context,
             "distinct_id": distinct_id,
             "$set_once": properties,
             "event": "$set_once",
@@ -304,13 +422,24 @@ class Client(object):
         timestamp=None,
         uuid=None,
         disable_geoip=None,
+        distinct_id=None,
     ):
+        if context is not None:
+            warnings.warn(
+                "The 'context' parameter is deprecated and will be removed in a future version.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
         properties = properties or {}
-        context = context or {}
         require("group_type", group_type, ID_TYPES)
         require("group_key", group_key, ID_TYPES)
         require("properties", properties, dict)
 
+        if distinct_id:
+            require("distinct_id", distinct_id, ID_TYPES)
+        else:
+            distinct_id = "${}_{}".format(group_type, group_key)
+
         msg = {
             "event": "$groupidentify",
             "properties": {
@@ -318,16 +447,20 @@ class Client(object):
                 "$group_key": group_key,
                 "$group_set": properties,
             },
-            "distinct_id": "${}_{}".format(group_type, group_key),
+            "distinct_id": distinct_id,
             "timestamp": timestamp,
-            "context": context,
             "uuid": uuid,
         }
 
         return self._enqueue(msg, disable_geoip)
 
     def alias(self, previous_id=None, distinct_id=None, context=None, timestamp=None, uuid=None, disable_geoip=None):
-        context = context or {}
+        if context is not None:
+            warnings.warn(
+                "The 'context' parameter is deprecated and will be removed in a future version.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
 
         require("previous_id", previous_id, ID_TYPES)
         require("distinct_id", distinct_id, ID_TYPES)
@@ -338,7 +471,6 @@ class Client(object):
                 "alias": distinct_id,
             },
             "timestamp": timestamp,
-            "context": context,
             "event": "$create_alias",
             "distinct_id": previous_id,
         }
@@ -348,9 +480,14 @@ class Client(object):
     def page(
         self, distinct_id=None, url=None, properties=None, context=None, timestamp=None, uuid=None, disable_geoip=None
     ):
-        properties = properties or {}
-        context = context or {}
+        if context is not None:
+            warnings.warn(
+                "The 'context' parameter is deprecated and will be removed in a future version.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
 
+        properties = properties or {}
         require("distinct_id", distinct_id, ID_TYPES)
         require("properties", properties, dict)
 
@@ -361,7 +498,6 @@ class Client(object):
             "event": "$pageview",
             "properties": properties,
             "timestamp": timestamp,
-            "context": context,
             "distinct_id": distinct_id,
             "uuid": uuid,
         }
@@ -378,6 +514,13 @@ class Client(object):
         uuid=None,
         groups=None,
     ):
+        if context is not None:
+            warnings.warn(
+                "The 'context' parameter is deprecated and will be removed in a future version.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
         # this function shouldn't ever throw an error, so it logs exceptions instead of raising them.
         # this is important to ensure we don't unexpectedly re-raise exceptions in the user's code.
         try:
@@ -438,7 +581,6 @@ class Client(object):
             timestamp = datetime.now(tz=tzutc())
 
         require("timestamp", timestamp, datetime)
-        require("context", msg["context"], dict)
 
         # add common
         timestamp = guess_timezone(timestamp)
@@ -535,9 +677,6 @@ class Client(object):
             )
 
             self.feature_flags = response["flags"] or []
-            self.feature_flags_by_key = {
-                flag["key"]: flag for flag in self.feature_flags if flag.get("key") is not None
-            }
             self.group_type_mapping = response["group_type_mapping"] or {}
             self.cohorts = response["cohorts"] or {}
 
@@ -553,6 +692,20 @@ 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.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:
@@ -584,7 +737,7 @@ class Client(object):
         person_properties={},
         group_properties={},
         warn_on_unknown_groups=True,
-    ):
+    ) -> FlagValue:
         if feature_flag.get("ensure_experience_continuity", False):
             raise InconclusiveMatchError("Flag has experience continuity enabled")
 
@@ -659,7 +812,14 @@ class Client(object):
         only_evaluate_locally=False,
         send_feature_flag_events=True,
         disable_geoip=None,
-    ):
+    ) -> Optional[FlagValue]:
+        """
+        Get a feature flag value for a key by evaluating locally or remotely
+        depending on whether local evaluation is enabled and the flag can be
+        locally evaluated.
+
+        This also captures the $feature_flag_called event unless send_feature_flag_events is False.
+        """
         require("key", key, string_types)
         require("distinct_id", distinct_id, ID_TYPES)
         require("groups", groups, dict)
@@ -671,65 +831,69 @@ class Client(object):
             distinct_id, groups, person_properties, group_properties
         )
 
-        if self.feature_flags is None and self.personal_api_key:
-            self.load_feature_flags()
-        response = None
+        response = self._locally_evaluate_flag(key, distinct_id, groups, person_properties, group_properties)
 
-        # If loading in previous line failed
-        if self.feature_flags:
-            for flag in self.feature_flags:
-                if flag["key"] == key:
-                    try:
-                        response = self._compute_flag_locally(
-                            flag,
-                            distinct_id,
-                            groups=groups,
-                            person_properties=person_properties,
-                            group_properties=group_properties,
-                        )
-                        self.log.debug(f"Successfully computed flag locally: {key} -> {response}")
-                    except InconclusiveMatchError as e:
-                        self.log.debug(f"Failed to compute flag {key} locally: {e}")
-                        continue
-                    except Exception as e:
-                        self.log.exception(f"[FEATURE FLAGS] Error while computing variant locally: {e}")
-                        continue
+        flag_details = None
+        request_id = None
 
         flag_was_locally_evaluated = response is not None
         if not flag_was_locally_evaluated and not only_evaluate_locally:
             try:
-                feature_flags = self.get_feature_variants(
-                    distinct_id,
-                    groups=groups,
-                    person_properties=person_properties,
-                    group_properties=group_properties,
-                    disable_geoip=disable_geoip,
+                flag_details, request_id = self._get_feature_flag_details_from_decide(
+                    key, distinct_id, groups, person_properties, group_properties, disable_geoip
                 )
-                response = feature_flags.get(key)
-                if response is None:
-                    response = False
+                response = flag_details.get_value() if flag_details else False
                 self.log.debug(f"Successfully computed flag remotely: #{key} -> #{response}")
             except Exception as e:
                 self.log.exception(f"[FEATURE FLAGS] Unable to get flag remotely: {e}")
 
-        feature_flag_reported_key = f"{key}_{str(response)}"
-        if (
-            feature_flag_reported_key not in self.distinct_ids_feature_flags_reported[distinct_id]
-            and send_feature_flag_events  # noqa: W503
-        ):
-            self.capture(
+        if send_feature_flag_events:
+            self._capture_feature_flag_called(
                 distinct_id,
-                "$feature_flag_called",
-                {
-                    "$feature_flag": key,
-                    "$feature_flag_response": response,
-                    "locally_evaluated": flag_was_locally_evaluated,
-                    f"$feature/{key}": response,
-                },
-                groups=groups,
-                disable_geoip=disable_geoip,
+                key,
+                response or False,
+                None,
+                flag_was_locally_evaluated,
+                groups,
+                disable_geoip,
+                request_id,
+                flag_details,
             )
-            self.distinct_ids_feature_flags_reported[distinct_id].add(feature_flag_reported_key)
+
+        return response
+
+    def _locally_evaluate_flag(
+        self,
+        key: str,
+        distinct_id: str,
+        groups: dict[str, str],
+        person_properties: dict[str, str],
+        group_properties: dict[str, str],
+    ) -> Optional[FlagValue]:
+        if self.feature_flags is None and self.personal_api_key:
+            self.load_feature_flags()
+        response = None
+
+        if self.feature_flags:
+            assert (
+                self.feature_flags_by_key is not None
+            ), "feature_flags_by_key should be initialized when feature_flags is set"
+            # Local evaluation
+            flag = self.feature_flags_by_key.get(key)
+            if flag:
+                try:
+                    response = self._compute_flag_locally(
+                        flag,
+                        distinct_id,
+                        groups=groups,
+                        person_properties=person_properties,
+                        group_properties=group_properties,
+                    )
+                    self.log.debug(f"Successfully computed flag locally: {key} -> {response}")
+                except InconclusiveMatchError as e:
+                    self.log.debug(f"Failed to compute flag {key} locally: {e}")
+                except Exception as e:
+                    self.log.exception(f"[FEATURE FLAGS] Error while computing variant locally: {e}")
         return response
 
     def get_feature_flag_payload(
@@ -749,21 +913,15 @@ class Client(object):
             return None
 
         if match_value is None:
-            match_value = self.get_feature_flag(
-                key,
-                distinct_id,
-                groups=groups,
-                person_properties=person_properties,
-                group_properties=group_properties,
-                send_feature_flag_events=False,
-                # Disable automatic sending of feature flag events because we're manually handling event dispatch.
-                # This prevents sending events with empty data when `get_feature_flag` cannot be evaluated locally.
-                only_evaluate_locally=True,  # Enable local evaluation of feature flags to avoid making multiple requests to `/decide`.
-                disable_geoip=disable_geoip,
+            person_properties, group_properties = self._add_local_person_and_group_properties(
+                distinct_id, groups, person_properties, group_properties
             )
+            match_value = self._locally_evaluate_flag(key, distinct_id, groups, person_properties, group_properties)
 
         response = None
         payload = None
+        flag_details = None
+        request_id = None
 
         if match_value is not None:
             payload = self._compute_payload_locally(key, match_value)
@@ -771,47 +929,126 @@ class Client(object):
         flag_was_locally_evaluated = payload is not None
         if not flag_was_locally_evaluated and not only_evaluate_locally:
             try:
-                responses_and_payloads = self.get_feature_flags_and_payloads(
-                    distinct_id, groups, person_properties, group_properties, disable_geoip
+                flag_details, request_id = self._get_feature_flag_details_from_decide(
+                    key, 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 = flag_details.metadata.payload if flag_details else None
+                response = flag_details.get_value() if flag_details else False
             except Exception as e:
                 self.log.exception(f"[FEATURE FLAGS] Unable to get feature flags and payloads: {e}")
 
+        if send_feature_flag_events:
+            self._capture_feature_flag_called(
+                distinct_id,
+                key,
+                response or False,
+                payload,
+                flag_was_locally_evaluated,
+                groups,
+                disable_geoip,
+                request_id,
+                flag_details,
+            )
+
+        return payload
+
+    def _get_feature_flag_details_from_decide(
+        self,
+        key: str,
+        distinct_id: str,
+        groups: dict[str, str],
+        person_properties: dict[str, str],
+        group_properties: dict[str, str],
+        disable_geoip: Optional[bool],
+    ) -> tuple[Optional[FeatureFlag], Optional[str]]:
+        """
+        Calls /decide and returns the flag details and request id
+        """
+        resp_data = self.get_decide(distinct_id, groups, person_properties, group_properties, disable_geoip)
+        request_id = resp_data.get("requestId")
+        flags = resp_data.get("flags")
+        flag_details = flags.get(key) if flags else None
+        return flag_details, request_id
+
+    def _capture_feature_flag_called(
+        self,
+        distinct_id: str,
+        key: str,
+        response: FlagValue,
+        payload: Optional[str],
+        flag_was_locally_evaluated: bool,
+        groups: dict[str, str],
+        disable_geoip: Optional[bool],
+        request_id: Optional[str],
+        flag_details: Optional[FeatureFlag],
+    ):
         feature_flag_reported_key = f"{key}_{str(response)}"
 
-        if (
-            feature_flag_reported_key not in self.distinct_ids_feature_flags_reported[distinct_id]
-            and send_feature_flag_events  # noqa: W503
-        ):
+        if feature_flag_reported_key not in self.distinct_ids_feature_flags_reported[distinct_id]:
+            properties: dict[str, Any] = {
+                "$feature_flag": key,
+                "$feature_flag_response": response,
+                "locally_evaluated": flag_was_locally_evaluated,
+                f"$feature/{key}": response,
+            }
+
+            if payload:
+                properties["$feature_flag_payload"] = payload
+
+            if request_id:
+                properties["$feature_flag_request_id"] = request_id
+            if isinstance(flag_details, FeatureFlag):
+                if flag_details.reason and flag_details.reason.description:
+                    properties["$feature_flag_reason"] = flag_details.reason.description
+                if isinstance(flag_details.metadata, FlagMetadata):
+                    if flag_details.metadata.version:
+                        properties["$feature_flag_version"] = flag_details.metadata.version
+                    if flag_details.metadata.id:
+                        properties["$feature_flag_id"] = flag_details.metadata.id
+
             self.capture(
                 distinct_id,
                 "$feature_flag_called",
-                {
-                    "$feature_flag": key,
-                    "$feature_flag_response": response,
-                    "$feature_flag_payload": payload,
-                    "locally_evaluated": flag_was_locally_evaluated,
-                    f"$feature/{key}": response,
-                },
+                properties,
                 groups=groups,
                 disable_geoip=disable_geoip,
             )
             self.distinct_ids_feature_flags_reported[distinct_id].add(feature_flag_reported_key)
 
-        return payload
+    def get_remote_config_payload(self, key: str):
+        if self.disabled:
+            return None
 
-    def _compute_payload_locally(self, key, match_value):
+        if self.personal_api_key is None:
+            self.log.warning(
+                "[FEATURE FLAGS] You have to specify a personal_api_key to fetch decrypted feature flag payloads."
+            )
+            return None
+
+        try:
+            return remote_config(
+                self.personal_api_key,
+                self.host,
+                key,
+                timeout=self.feature_flags_request_timeout_seconds,
+            )
+        except Exception as e:
+            self.log.exception(f"[FEATURE FLAGS] Unable to get decrypted feature flag payload: {e}")
+
+    def _compute_payload_locally(self, key: str, match_value: FlagValue) -> Optional[str]:
         payload = None
 
         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(
@@ -823,8 +1060,8 @@ class Client(object):
         group_properties={},
         only_evaluate_locally=False,
         disable_geoip=None,
-    ):
-        flags = self.get_all_flags_and_payloads(
+    ) -> Optional[dict[str, Union[bool, str]]]:
+        response = self.get_all_flags_and_payloads(
             distinct_id,
             groups=groups,
             person_properties=person_properties,
@@ -832,7 +1069,8 @@ class Client(object):
             only_evaluate_locally=only_evaluate_locally,
             disable_geoip=disable_geoip,
         )
-        return flags["featureFlags"]
+
+        return response["featureFlags"]
 
     def get_all_flags_and_payloads(
         self,
@@ -843,7 +1081,7 @@ class Client(object):
         group_properties={},
         only_evaluate_locally=False,
         disable_geoip=None,
-    ):
+    ) -> FlagsAndPayloads:
         if self.disabled:
             return {"featureFlags": None, "featureFlagPayloads": None}
 
@@ -851,21 +1089,20 @@ class Client(object):
             distinct_id, groups, person_properties, group_properties
         )
 
-        flags, payloads, fallback_to_decide = self._get_all_flags_and_payloads_locally(
+        response, fallback_to_decide = self._get_all_flags_and_payloads_locally(
             distinct_id, groups=groups, person_properties=person_properties, group_properties=group_properties
         )
-        response = {"featureFlags": flags, "featureFlagPayloads": payloads}
 
         if fallback_to_decide and not only_evaluate_locally:
             try:
-                flags_and_payloads = self.get_decide(
+                decide_response = self.get_decide(
                     distinct_id,
                     groups=groups,
                     person_properties=person_properties,
                     group_properties=group_properties,
                     disable_geoip=disable_geoip,
                 )
-                response = flags_and_payloads
+                return to_flags_and_payloads(decide_response)
             except Exception as e:
                 self.log.exception(f"[FEATURE FLAGS] Unable to get feature flags and payloads: {e}")
 
@@ -873,15 +1110,15 @@ class Client(object):
 
     def _get_all_flags_and_payloads_locally(
         self, distinct_id, *, groups={}, person_properties={}, group_properties={}, warn_on_unknown_groups=False
-    ):
+    ) -> tuple[FlagsAndPayloads, bool]:
         require("distinct_id", distinct_id, ID_TYPES)
         require("groups", groups, dict)
 
         if self.feature_flags is None and self.personal_api_key:
             self.load_feature_flags()
 
-        flags = {}
-        payloads = {}
+        flags: dict[str, FlagValue] = {}
+        payloads: dict[str, str] = {}
         fallback_to_decide = False
         # If loading in previous line failed
         if self.feature_flags:
@@ -907,7 +1144,7 @@ class Client(object):
         else:
             fallback_to_decide = True
 
-        return flags, payloads, fallback_to_decide
+        return {"featureFlags": flags, "featureFlagPayloads": payloads}, fallback_to_decide
 
     def feature_flag_definitions(self):
         return self.feature_flags

posthog/exception_utils.py

@@ -793,7 +793,7 @@ def event_from_exception(
 
 
 def _module_in_list(name, items):
-    # type: (str, Optional[List[str]]) -> bool
+    # type: (str | None, Optional[List[str]]) -> bool
     if name is None:
         return False
 

posthog/feature_flags.py

@@ -7,6 +7,8 @@ from typing import Optional
 from dateutil import parser
 from dateutil.relativedelta import relativedelta
 
+from posthog import utils
+from posthog.types import FlagValue
 from posthog.utils import convert_to_datetime_aware, is_valid_regex
 
 __LONG_SCALE__ = float(0xFFFFFFFFFFFFFFF)
@@ -24,7 +26,7 @@ class InconclusiveMatchError(Exception):
 # Given the same distinct_id and key, it'll always return the same float. These floats are
 # uniformly distributed between 0 and 1, so if we want to show this feature to 20% of traffic
 # we can do _hash(key, distinct_id) < 0.2
-def _hash(key, distinct_id, salt=""):
+def _hash(key: str, distinct_id: str, salt: str = "") -> float:
     hash_key = f"{key}.{distinct_id}{salt}"
     hash_val = int(hashlib.sha1(hash_key.encode("utf-8")).hexdigest()[:15], 16)
     return hash_val / __LONG_SCALE__
@@ -49,10 +51,13 @@ def variant_lookup_table(feature_flag):
     return lookup_table
 
 
-def match_feature_flag_properties(flag, distinct_id, properties, cohort_properties=None):
+def match_feature_flag_properties(flag, distinct_id, properties, cohort_properties=None) -> FlagValue:
     flag_conditions = (flag.get("filters") or {}).get("groups") or []
     is_inconclusive = False
     cohort_properties = cohort_properties or {}
+    # Some filters can be explicitly set to null, which require accessing variants like so
+    flag_variants = ((flag.get("filters") or {}).get("multivariate") or {}).get("variants") or []
+    valid_variant_keys = [variant["key"] for variant in flag_variants]
 
     # Stable sort conditions with variant overrides to the top. This ensures that if overrides are present, they are
     # evaluated first, and the variant override is applied to the first matching condition.
@@ -67,9 +72,7 @@ def match_feature_flag_properties(flag, distinct_id, properties, cohort_properti
             # the matching variant
             if is_condition_match(flag, distinct_id, condition, properties, cohort_properties):
                 variant_override = condition.get("variant")
-                # Some filters can be explicitly set to null, which require accessing variants like so
-                flag_variants = ((flag.get("filters") or {}).get("multivariate") or {}).get("variants") or []
-                if variant_override and variant_override in [variant["key"] for variant in flag_variants]:
+                if variant_override and variant_override in valid_variant_keys:
                     variant = variant_override
                 else:
                     variant = get_matching_variant(flag, distinct_id)
@@ -85,7 +88,7 @@ def match_feature_flag_properties(flag, distinct_id, properties, cohort_properti
     return False
 
 
-def is_condition_match(feature_flag, distinct_id, condition, properties, cohort_properties):
+def is_condition_match(feature_flag, distinct_id, condition, properties, cohort_properties) -> bool:
     rollout_percentage = condition.get("rollout_percentage")
     if len(condition.get("properties") or []) > 0:
         for prop in condition.get("properties"):
@@ -128,8 +131,8 @@ def match_property(property, property_values) -> bool:
 
         def compute_exact_match(value, override_value):
             if isinstance(value, list):
-                return str(override_value).lower() in [str(val).lower() for val in value]
-            return str(value).lower() == str(override_value).lower()
+                return str(override_value).casefold() in [str(val).casefold() for val in value]
+            return utils.str_iequals(value, override_value)
 
         if operator == "exact":
             return compute_exact_match(value, override_value)
@@ -140,10 +143,10 @@ def match_property(property, property_values) -> bool:
         return key in property_values
 
     if operator == "icontains":
-        return str(value).lower() in str(override_value).lower()
+        return utils.str_icontains(override_value, value)
 
     if operator == "not_icontains":
-        return str(value).lower() not in str(override_value).lower()
+        return not utils.str_icontains(override_value, value)
 
     if operator == "regex":
         return is_valid_regex(str(value)) and re.compile(str(value)).search(str(override_value)) is not None

posthog/request.py

@@ -7,11 +7,22 @@ from typing import Any, Optional, Union
 
 import requests
 from dateutil.tz import tzutc
+from urllib3.util.retry import Retry
 
 from posthog.utils import remove_trailing_slash
 from posthog.version import VERSION
 
+# Retry on both connect and read errors
+# by default read errors will only retry idempotent HTTP methods (so not POST)
+adapter = requests.adapters.HTTPAdapter(
+    max_retries=Retry(
+        total=2,
+        connect=2,
+        read=2,
+    )
+)
 _session = requests.sessions.Session()
+_session.mount("https://", adapter)
 
 US_INGESTION_ENDPOINT = "https://us.i.posthog.com"
 EU_INGESTION_ENDPOINT = "https://eu.i.posthog.com"
@@ -66,7 +77,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)
@@ -77,10 +102,15 @@ def _process_response(
 
 def decide(api_key: str, host: Optional[str] = None, gzip: bool = False, timeout: int = 15, **kwargs) -> Any:
     """Post the `kwargs to the decide API endpoint"""
-    res = post(api_key, host, "/decide/?v=3", gzip, timeout, **kwargs)
+    res = post(api_key, host, "/decide/?v=4", gzip, timeout, **kwargs)
     return _process_response(res, success_message="Feature flags decided successfully")
 
 
+def remote_config(personal_api_key: str, host: Optional[str] = None, key: str = "", timeout: int = 15) -> Any:
+    """Get remote config flag value from remote_config API endpoint"""
+    return get(personal_api_key, f"/api/projects/@current/feature_flags/{key}/remote_config/", host, timeout)
+
+
 def batch_post(
     api_key: str, host: Optional[str] = None, gzip: bool = False, timeout: int = 15, **kwargs
 ) -> requests.Response:
@@ -105,6 +135,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)):

posthog/test/ai/anthropic/test_anthropic.py

@@ -0,0 +1,394 @@
+import os
+import time
+from unittest.mock import patch
+
+import pytest
+from anthropic.types import Message, Usage
+
+from posthog.ai.anthropic import Anthropic, AsyncAnthropic
+
+ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY")
+
+
+@pytest.fixture
+def mock_client():
+    with patch("posthog.client.Client") as mock_client:
+        mock_client.privacy_mode = False
+        yield mock_client
+
+
+@pytest.fixture
+def mock_anthropic_response():
+    return Message(
+        id="msg_123",
+        type="message",
+        role="assistant",
+        content=[{"type": "text", "text": "Test response"}],
+        model="claude-3-opus-20240229",
+        usage=Usage(
+            input_tokens=20,
+            output_tokens=10,
+        ),
+        stop_reason="end_turn",
+        stop_sequence=None,
+    )
+
+
+@pytest.fixture
+def mock_anthropic_stream():
+    class MockStreamEvent:
+        def __init__(self, content, usage=None):
+            self.content = content
+            self.usage = usage
+
+    def stream_generator():
+        yield MockStreamEvent("A")
+        yield MockStreamEvent("B")
+        yield MockStreamEvent(
+            "C",
+            usage=Usage(
+                input_tokens=20,
+                output_tokens=10,
+            ),
+        )
+
+    return stream_generator()
+
+
+@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)
+        response = client.messages.create(
+            model="claude-3-opus-20240229",
+            messages=[{"role": "user", "content": "Hello"}],
+            posthog_distinct_id="test-id",
+            posthog_properties={"foo": "bar"},
+        )
+
+        assert response == mock_anthropic_response
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+
+        assert call_args["distinct_id"] == "test-id"
+        assert call_args["event"] == "$ai_generation"
+        assert props["$ai_provider"] == "anthropic"
+        assert props["$ai_model"] == "claude-3-opus-20240229"
+        assert props["$ai_input"] == [{"role": "user", "content": "Hello"}]
+        assert props["$ai_output_choices"] == [{"role": "assistant", "content": "Test response"}]
+        assert props["$ai_input_tokens"] == 20
+        assert props["$ai_output_tokens"] == 10
+        assert props["$ai_http_status"] == 200
+        assert props["foo"] == "bar"
+        assert isinstance(props["$ai_latency"], float)
+
+
+def test_streaming(mock_client, mock_anthropic_stream):
+    with patch("anthropic.resources.Messages.create", return_value=mock_anthropic_stream):
+        client = Anthropic(api_key="test-key", posthog_client=mock_client)
+        response = client.messages.create(
+            model="claude-3-opus-20240229",
+            messages=[{"role": "user", "content": "Hello"}],
+            stream=True,
+            posthog_distinct_id="test-id",
+            posthog_properties={"foo": "bar"},
+        )
+
+        # Consume the stream
+        chunks = list(response)
+        assert len(chunks) == 3
+        assert chunks[0].content == "A"
+        assert chunks[1].content == "B"
+        assert chunks[2].content == "C"
+
+        # Wait a bit to ensure the capture is called
+        time.sleep(0.1)
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+
+        assert call_args["distinct_id"] == "test-id"
+        assert call_args["event"] == "$ai_generation"
+        assert props["$ai_provider"] == "anthropic"
+        assert props["$ai_model"] == "claude-3-opus-20240229"
+        assert props["$ai_input"] == [{"role": "user", "content": "Hello"}]
+        assert props["$ai_output_choices"] == [{"role": "assistant", "content": "ABC"}]
+        assert props["$ai_input_tokens"] == 20
+        assert props["$ai_output_tokens"] == 10
+        assert isinstance(props["$ai_latency"], float)
+        assert props["foo"] == "bar"
+
+
+def test_streaming_with_stream_endpoint(mock_client, mock_anthropic_stream):
+    with patch("anthropic.resources.Messages.create", return_value=mock_anthropic_stream):
+        client = Anthropic(api_key="test-key", posthog_client=mock_client)
+        response = client.messages.stream(
+            model="claude-3-opus-20240229",
+            messages=[{"role": "user", "content": "Hello"}],
+            posthog_distinct_id="test-id",
+            posthog_properties={"foo": "bar"},
+        )
+
+        # Consume the stream
+        chunks = list(response)
+        assert len(chunks) == 3
+        assert chunks[0].content == "A"
+        assert chunks[1].content == "B"
+        assert chunks[2].content == "C"
+
+        # Wait a bit to ensure the capture is called
+        time.sleep(0.1)
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+
+        assert call_args["distinct_id"] == "test-id"
+        assert call_args["event"] == "$ai_generation"
+        assert props["$ai_provider"] == "anthropic"
+        assert props["$ai_model"] == "claude-3-opus-20240229"
+        assert props["$ai_input"] == [{"role": "user", "content": "Hello"}]
+        assert props["$ai_output_choices"] == [{"role": "assistant", "content": "ABC"}]
+        assert props["$ai_input_tokens"] == 20
+        assert props["$ai_output_tokens"] == 10
+        assert isinstance(props["$ai_latency"], float)
+        assert props["foo"] == "bar"
+
+
+def test_groups(mock_client, mock_anthropic_response):
+    with patch("anthropic.resources.Messages.create", return_value=mock_anthropic_response):
+        client = Anthropic(api_key="test-key", posthog_client=mock_client)
+        response = client.messages.create(
+            model="claude-3-opus-20240229",
+            messages=[{"role": "user", "content": "Hello"}],
+            posthog_distinct_id="test-id",
+            posthog_groups={"company": "test_company"},
+        )
+
+        assert response == mock_anthropic_response
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+        assert call_args["groups"] == {"company": "test_company"}
+
+
+def test_privacy_mode_local(mock_client, mock_anthropic_response):
+    with patch("anthropic.resources.Messages.create", return_value=mock_anthropic_response):
+        client = Anthropic(api_key="test-key", posthog_client=mock_client)
+        response = client.messages.create(
+            model="claude-3-opus-20240229",
+            messages=[{"role": "user", "content": "Hello"}],
+            posthog_distinct_id="test-id",
+            posthog_privacy_mode=True,
+        )
+
+        assert response == mock_anthropic_response
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+        assert props["$ai_input"] is None
+        assert props["$ai_output_choices"] is None
+
+
+def test_privacy_mode_global(mock_client, mock_anthropic_response):
+    with patch("anthropic.resources.Messages.create", return_value=mock_anthropic_response):
+        mock_client.privacy_mode = True
+        client = Anthropic(api_key="test-key", posthog_client=mock_client)
+        response = client.messages.create(
+            model="claude-3-opus-20240229",
+            messages=[{"role": "user", "content": "Hello"}],
+            posthog_distinct_id="test-id",
+            posthog_privacy_mode=False,
+        )
+
+        assert response == mock_anthropic_response
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+        assert props["$ai_input"] is None
+        assert props["$ai_output_choices"] is None
+
+
+@pytest.mark.skipif(not ANTHROPIC_API_KEY, reason="ANTHROPIC_API_KEY is not set")
+def test_basic_integration(mock_client):
+    client = Anthropic(posthog_client=mock_client)
+    client.messages.create(
+        model="claude-3-opus-20240229",
+        messages=[{"role": "user", "content": "Foo"}],
+        max_tokens=1,
+        temperature=0,
+        posthog_distinct_id="test-id",
+        posthog_properties={"foo": "bar"},
+        system="You must always answer with 'Bar'.",
+    )
+
+    assert mock_client.capture.call_count == 1
+
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+    assert call_args["distinct_id"] == "test-id"
+    assert call_args["event"] == "$ai_generation"
+    assert props["$ai_provider"] == "anthropic"
+    assert props["$ai_model"] == "claude-3-opus-20240229"
+    assert props["$ai_input"] == [
+        {"role": "system", "content": "You must always answer with 'Bar'."},
+        {"role": "user", "content": "Foo"},
+    ]
+    assert props["$ai_output_choices"][0]["role"] == "assistant"
+    assert props["$ai_output_choices"][0]["content"] == "Bar"
+    assert props["$ai_input_tokens"] == 18
+    assert props["$ai_output_tokens"] == 1
+    assert props["$ai_http_status"] == 200
+    assert props["foo"] == "bar"
+    assert isinstance(props["$ai_latency"], float)
+
+
+@pytest.mark.skipif(not ANTHROPIC_API_KEY, reason="ANTHROPIC_API_KEY is not set")
+async def test_basic_async_integration(mock_client):
+    client = AsyncAnthropic(posthog_client=mock_client)
+    await client.messages.create(
+        model="claude-3-opus-20240229",
+        messages=[{"role": "user", "content": "You must always answer with 'Bar'."}],
+        max_tokens=1,
+        temperature=0,
+        posthog_distinct_id="test-id",
+        posthog_properties={"foo": "bar"},
+    )
+
+    assert mock_client.capture.call_count == 1
+
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+
+    assert call_args["distinct_id"] == "test-id"
+    assert call_args["event"] == "$ai_generation"
+    assert props["$ai_provider"] == "anthropic"
+    assert props["$ai_model"] == "claude-3-opus-20240229"
+    assert props["$ai_input"] == [{"role": "user", "content": "You must always answer with 'Bar'."}]
+    assert props["$ai_output_choices"][0]["role"] == "assistant"
+    assert props["$ai_input_tokens"] == 16
+    assert props["$ai_output_tokens"] == 1
+    assert props["$ai_http_status"] == 200
+    assert props["foo"] == "bar"
+    assert isinstance(props["$ai_latency"], float)
+
+
+def test_streaming_system_prompt(mock_client, mock_anthropic_stream):
+    with patch("anthropic.resources.Messages.create", return_value=mock_anthropic_stream):
+        client = Anthropic(api_key="test-key", posthog_client=mock_client)
+        response = client.messages.create(
+            model="claude-3-opus-20240229",
+            system="Foo",
+            messages=[{"role": "user", "content": "Bar"}],
+            stream=True,
+        )
+
+        # Consume the stream
+        list(response)
+
+        # Wait a bit to ensure the capture is called
+        time.sleep(0.1)
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+
+        assert props["$ai_input"] == [{"role": "system", "content": "Foo"}, {"role": "user", "content": "Bar"}]
+
+
+@pytest.mark.skipif(not ANTHROPIC_API_KEY, reason="ANTHROPIC_API_KEY is not set")
+async def test_async_streaming_system_prompt(mock_client, mock_anthropic_stream):
+    client = AsyncAnthropic(posthog_client=mock_client)
+    response = await client.messages.create(
+        model="claude-3-opus-20240229",
+        system="You must always answer with 'Bar'.",
+        messages=[{"role": "user", "content": "Foo"}],
+        stream=True,
+        max_tokens=1,
+    )
+
+    # Consume the stream
+    [c async for c in response]
+
+    # Wait a bit to ensure the capture is called
+    time.sleep(0.1)
+    assert mock_client.capture.call_count == 1
+
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+
+    assert props["$ai_input"] == [
+        {"role": "system", "content": "You must always answer with 'Bar'."},
+        {"role": "user", "content": "Foo"},
+    ]
+
+
+def test_error(mock_client, mock_anthropic_response):
+    with patch("anthropic.resources.Messages.create", side_effect=Exception("Test error")):
+        client = Anthropic(api_key="test-key", posthog_client=mock_client)
+        with pytest.raises(Exception):
+            client.messages.create(model="claude-3-opus-20240229", messages=[{"role": "user", "content": "Hello"}])
+
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+        assert props["$ai_is_error"] is True
+        assert props["$ai_error"] == "Test error"
+
+
+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)

posthog/test/ai/langchain/__init__.py

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

posthog/test/ai/openai/test_openai.py

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

posthog/test/test_client.py

@@ -5,9 +5,12 @@ 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.types import FeatureFlag, LegacyFlagMetadata
 from posthog.version import VERSION
 
 
@@ -53,6 +56,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 +108,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 +135,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 +162,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 +189,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 +216,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 +252,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()
 
@@ -292,7 +293,6 @@ class TestClient(unittest.TestCase):
             "id": 1,
             "name": "Beta Feature",
             "key": "beta-feature-local",
-            "is_simple_flag": False,
             "active": True,
             "rollout_percentage": 100,
             "filters": {
@@ -321,7 +321,6 @@ class TestClient(unittest.TestCase):
             "id": 1,
             "name": "Beta Feature",
             "key": "person-flag",
-            "is_simple_flag": True,
             "active": True,
             "filters": {
                 "groups": [
@@ -344,7 +343,6 @@ class TestClient(unittest.TestCase):
             "id": 1,
             "name": "Beta Feature",
             "key": "false-flag",
-            "is_simple_flag": True,
             "active": True,
             "filters": {
                 "groups": [
@@ -387,6 +385,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"}}
@@ -396,7 +413,6 @@ class TestClient(unittest.TestCase):
             "id": 1,
             "name": "Beta Feature",
             "key": "beta-feature-local",
-            "is_simple_flag": False,
             "active": True,
             "rollout_percentage": 100,
             "filters": {
@@ -425,7 +441,6 @@ class TestClient(unittest.TestCase):
             "id": 1,
             "name": "Beta Feature",
             "key": "person-flag",
-            "is_simple_flag": True,
             "active": True,
             "filters": {
                 "groups": [
@@ -581,16 +596,14 @@ class TestClient(unittest.TestCase):
             "distinct_id",
             "python test event",
             {"property": "value"},
-            {"ip": "192.168.0.1"},
-            datetime(2014, 9, 3),
-            "new-uuid",
+            timestamp=datetime(2014, 9, 3),
+            uuid="new-uuid",
         )
 
         self.assertTrue(success)
 
         self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00")
         self.assertEqual(msg["properties"]["property"], "value")
-        self.assertEqual(msg["context"]["ip"], "192.168.0.1")
         self.assertEqual(msg["event"], "python test event")
         self.assertEqual(msg["properties"]["$lib"], "posthog-python")
         self.assertEqual(msg["properties"]["$lib_version"], VERSION)
@@ -623,13 +636,12 @@ class TestClient(unittest.TestCase):
     def test_advanced_identify(self):
         client = self.client
         success, msg = client.identify(
-            "distinct_id", {"trait": "value"}, {"ip": "192.168.0.1"}, datetime(2014, 9, 3), "new-uuid"
+            "distinct_id", {"trait": "value"}, timestamp=datetime(2014, 9, 3), uuid="new-uuid"
         )
 
         self.assertTrue(success)
 
         self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00")
-        self.assertEqual(msg["context"]["ip"], "192.168.0.1")
         self.assertEqual(msg["$set"]["trait"], "value")
         self.assertEqual(msg["properties"]["$lib"], "posthog-python")
         self.assertEqual(msg["properties"]["$lib_version"], VERSION)
@@ -651,14 +663,11 @@ class TestClient(unittest.TestCase):
 
     def test_advanced_set(self):
         client = self.client
-        success, msg = client.set(
-            "distinct_id", {"trait": "value"}, {"ip": "192.168.0.1"}, datetime(2014, 9, 3), "new-uuid"
-        )
+        success, msg = client.set("distinct_id", {"trait": "value"}, timestamp=datetime(2014, 9, 3), uuid="new-uuid")
 
         self.assertTrue(success)
 
         self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00")
-        self.assertEqual(msg["context"]["ip"], "192.168.0.1")
         self.assertEqual(msg["$set"]["trait"], "value")
         self.assertEqual(msg["properties"]["$lib"], "posthog-python")
         self.assertEqual(msg["properties"]["$lib_version"], VERSION)
@@ -681,13 +690,12 @@ class TestClient(unittest.TestCase):
     def test_advanced_set_once(self):
         client = self.client
         success, msg = client.set_once(
-            "distinct_id", {"trait": "value"}, {"ip": "192.168.0.1"}, datetime(2014, 9, 3), "new-uuid"
+            "distinct_id", {"trait": "value"}, timestamp=datetime(2014, 9, 3), uuid="new-uuid"
         )
 
         self.assertTrue(success)
 
         self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00")
-        self.assertEqual(msg["context"]["ip"], "192.168.0.1")
         self.assertEqual(msg["$set_once"]["trait"], "value")
         self.assertEqual(msg["properties"]["$lib"], "posthog-python")
         self.assertEqual(msg["properties"]["$lib_version"], VERSION)
@@ -715,9 +723,28 @@ class TestClient(unittest.TestCase):
         self.assertTrue(isinstance(msg["timestamp"], str))
         self.assertIsNone(msg.get("uuid"))
 
+    def test_basic_group_identify_with_distinct_id(self):
+        success, msg = self.client.group_identify("organization", "id:5", distinct_id="distinct_id")
+        self.assertTrue(success)
+        self.assertEqual(msg["event"], "$groupidentify")
+        self.assertEqual(msg["distinct_id"], "distinct_id")
+        self.assertEqual(
+            msg["properties"],
+            {
+                "$group_type": "organization",
+                "$group_key": "id:5",
+                "$group_set": {},
+                "$lib": "posthog-python",
+                "$lib_version": VERSION,
+                "$geoip_disable": True,
+            },
+        )
+        self.assertTrue(isinstance(msg["timestamp"], str))
+        self.assertIsNone(msg.get("uuid"))
+
     def test_advanced_group_identify(self):
         success, msg = self.client.group_identify(
-            "organization", "id:5", {"trait": "value"}, {"ip": "192.168.0.1"}, datetime(2014, 9, 3), "new-uuid"
+            "organization", "id:5", {"trait": "value"}, timestamp=datetime(2014, 9, 3), uuid="new-uuid"
         )
 
         self.assertTrue(success)
@@ -735,7 +762,33 @@ class TestClient(unittest.TestCase):
             },
         )
         self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00")
-        self.assertEqual(msg["context"]["ip"], "192.168.0.1")
+
+    def test_advanced_group_identify_with_distinct_id(self):
+        success, msg = self.client.group_identify(
+            "organization",
+            "id:5",
+            {"trait": "value"},
+            timestamp=datetime(2014, 9, 3),
+            uuid="new-uuid",
+            distinct_id="distinct_id",
+        )
+
+        self.assertTrue(success)
+        self.assertEqual(msg["event"], "$groupidentify")
+        self.assertEqual(msg["distinct_id"], "distinct_id")
+
+        self.assertEqual(
+            msg["properties"],
+            {
+                "$group_type": "organization",
+                "$group_key": "id:5",
+                "$group_set": {"trait": "value"},
+                "$lib": "posthog-python",
+                "$lib_version": VERSION,
+                "$geoip_disable": True,
+            },
+        )
+        self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00")
 
     def test_basic_alias(self):
         client = self.client
@@ -771,15 +824,13 @@ class TestClient(unittest.TestCase):
             "distinct_id",
             "https://posthog.com/contact",
             {"property": "value"},
-            {"ip": "192.168.0.1"},
-            datetime(2014, 9, 3),
-            "new-uuid",
+            timestamp=datetime(2014, 9, 3),
+            uuid="new-uuid",
         )
 
         self.assertTrue(success)
 
         self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00")
-        self.assertEqual(msg["context"]["ip"], "192.168.0.1")
         self.assertEqual(msg["properties"]["$current_url"], "https://posthog.com/contact")
         self.assertEqual(msg["properties"]["property"], "value")
         self.assertEqual(msg["properties"]["$lib"], "posthog-python")
@@ -1003,7 +1054,7 @@ class TestClient(unittest.TestCase):
 
         patch_get.return_value.raiseError.side_effect = raise_effect
         client = Client(FAKE_TEST_API_KEY, personal_api_key="test")
-        client.feature_flags = [{"key": "example", "is_simple_flag": False}]
+        client.feature_flags = [{"key": "example"}]
 
         self.assertFalse(client.feature_enabled("example", "distinct_id"))
 
@@ -1073,3 +1124,142 @@ 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
+
+    @mock.patch("posthog.client.decide")
+    def test_get_decide_returns_normalized_decide_response(self, patch_decide):
+        patch_decide.return_value = {
+            "featureFlags": {"beta-feature": "random-variant", "alpha-feature": True, "off-feature": False},
+            "featureFlagPayloads": {"beta-feature": '{"some": "data"}'},
+            "errorsWhileComputingFlags": False,
+            "requestId": "test-id",
+        }
+
+        client = Client(FAKE_TEST_API_KEY)
+        distinct_id = "test_distinct_id"
+        groups = {"test_group_type": "test_group_id"}
+        person_properties = {"test_property": "test_value"}
+
+        response = client.get_decide(distinct_id, groups, person_properties)
+
+        assert response == {
+            "flags": {
+                "beta-feature": FeatureFlag(
+                    key="beta-feature",
+                    enabled=True,
+                    variant="random-variant",
+                    reason=None,
+                    metadata=LegacyFlagMetadata(
+                        payload='{"some": "data"}',
+                    ),
+                ),
+                "alpha-feature": FeatureFlag(
+                    key="alpha-feature",
+                    enabled=True,
+                    variant=None,
+                    reason=None,
+                    metadata=LegacyFlagMetadata(
+                        payload=None,
+                    ),
+                ),
+                "off-feature": FeatureFlag(
+                    key="off-feature",
+                    enabled=False,
+                    variant=None,
+                    reason=None,
+                    metadata=LegacyFlagMetadata(
+                        payload=None,
+                    ),
+                ),
+            },
+            "errorsWhileComputingFlags": False,
+            "requestId": "test-id",
+        }

posthog/test/test_feature_flag.py

@@ -0,0 +1,142 @@
+import unittest
+
+from posthog.types import FeatureFlag, FlagMetadata, FlagReason, LegacyFlagMetadata
+
+
+class TestFeatureFlag(unittest.TestCase):
+    def test_feature_flag_from_json(self):
+        # Test with full metadata
+        resp = {
+            "key": "test-flag",
+            "enabled": True,
+            "variant": "test-variant",
+            "reason": {"code": "matched_condition", "condition_index": 0, "description": "Matched condition set 1"},
+            "metadata": {"id": 1, "payload": '{"some": "json"}', "version": 2, "description": "test-description"},
+        }
+
+        flag = FeatureFlag.from_json(resp)
+        self.assertEqual(flag.key, "test-flag")
+        self.assertTrue(flag.enabled)
+        self.assertEqual(flag.variant, "test-variant")
+        self.assertEqual(flag.get_value(), "test-variant")
+        self.assertEqual(
+            flag.reason, FlagReason(code="matched_condition", condition_index=0, description="Matched condition set 1")
+        )
+        self.assertEqual(
+            flag.metadata, FlagMetadata(id=1, payload='{"some": "json"}', version=2, description="test-description")
+        )
+
+    def test_feature_flag_from_json_minimal(self):
+        # Test with minimal required fields
+        resp = {"key": "test-flag", "enabled": True}
+
+        flag = FeatureFlag.from_json(resp)
+        self.assertEqual(flag.key, "test-flag")
+        self.assertTrue(flag.enabled)
+        self.assertIsNone(flag.variant)
+        self.assertEqual(flag.get_value(), True)
+        self.assertIsNone(flag.reason)
+        self.assertEqual(flag.metadata, LegacyFlagMetadata(payload=None))
+
+    def test_feature_flag_from_json_without_metadata(self):
+        # Test with reason but no metadata
+        resp = {
+            "key": "test-flag",
+            "enabled": True,
+            "variant": "test-variant",
+            "reason": {"code": "matched_condition", "condition_index": 0, "description": "Matched condition set 1"},
+        }
+
+        flag = FeatureFlag.from_json(resp)
+        self.assertEqual(flag.key, "test-flag")
+        self.assertTrue(flag.enabled)
+        self.assertEqual(flag.variant, "test-variant")
+        self.assertEqual(flag.get_value(), "test-variant")
+        self.assertEqual(
+            flag.reason, FlagReason(code="matched_condition", condition_index=0, description="Matched condition set 1")
+        )
+        self.assertEqual(flag.metadata, LegacyFlagMetadata(payload=None))
+
+    def test_flag_reason_from_json(self):
+        # Test with complete data
+        resp = {"code": "user_in_segment", "condition_index": 1, "description": "User is in segment 'beta_users'"}
+        reason = FlagReason.from_json(resp)
+        self.assertEqual(reason.code, "user_in_segment")
+        self.assertEqual(reason.condition_index, 1)
+        self.assertEqual(reason.description, "User is in segment 'beta_users'")
+
+        # Test with partial data
+        resp = {"code": "user_in_segment"}
+        reason = FlagReason.from_json(resp)
+        self.assertEqual(reason.code, "user_in_segment")
+        self.assertEqual(reason.condition_index, 0)  # default value
+        self.assertEqual(reason.description, "")  # default value
+
+        # Test with None
+        self.assertIsNone(FlagReason.from_json(None))
+
+    def test_flag_metadata_from_json(self):
+        # Test with complete data
+        resp = {"id": 123, "payload": {"key": "value"}, "version": 1, "description": "Test flag"}
+        metadata = FlagMetadata.from_json(resp)
+        self.assertEqual(metadata.id, 123)
+        self.assertEqual(metadata.payload, {"key": "value"})
+        self.assertEqual(metadata.version, 1)
+        self.assertEqual(metadata.description, "Test flag")
+
+        # Test with partial data
+        resp = {"id": 123}
+        metadata = FlagMetadata.from_json(resp)
+        self.assertEqual(metadata.id, 123)
+        self.assertIsNone(metadata.payload)
+        self.assertEqual(metadata.version, 0)  # default value
+        self.assertEqual(metadata.description, "")  # default value
+
+        # Test with None
+        self.assertIsInstance(FlagMetadata.from_json(None), LegacyFlagMetadata)
+
+    def test_feature_flag_from_json_complete(self):
+        # Test with complete data
+        resp = {
+            "key": "test-flag",
+            "enabled": True,
+            "variant": "control",
+            "reason": {
+                "code": "user_in_segment",
+                "condition_index": 1,
+                "description": "User is in segment 'beta_users'",
+            },
+            "metadata": {"id": 123, "payload": {"key": "value"}, "version": 1, "description": "Test flag"},
+        }
+        flag = FeatureFlag.from_json(resp)
+        self.assertEqual(flag.key, "test-flag")
+        self.assertTrue(flag.enabled)
+        self.assertEqual(flag.variant, "control")
+        self.assertIsInstance(flag.reason, FlagReason)
+        self.assertEqual(flag.reason.code, "user_in_segment")
+        self.assertIsInstance(flag.metadata, FlagMetadata)
+        self.assertEqual(flag.metadata.id, 123)
+        self.assertEqual(flag.metadata.payload, {"key": "value"})
+
+    def test_feature_flag_from_json_minimal_data(self):
+        # Test with minimal data
+        resp = {"key": "test-flag", "enabled": False}
+        flag = FeatureFlag.from_json(resp)
+        self.assertEqual(flag.key, "test-flag")
+        self.assertFalse(flag.enabled)
+        self.assertIsNone(flag.variant)
+        self.assertIsNone(flag.reason)
+        self.assertIsInstance(flag.metadata, LegacyFlagMetadata)
+        self.assertIsNone(flag.metadata.payload)
+
+    def test_feature_flag_from_json_with_reason(self):
+        # Test with reason but no metadata
+        resp = {"key": "test-flag", "enabled": True, "reason": {"code": "user_in_segment"}}
+        flag = FeatureFlag.from_json(resp)
+        self.assertEqual(flag.key, "test-flag")
+        self.assertTrue(flag.enabled)
+        self.assertIsNone(flag.variant)
+        self.assertIsInstance(flag.reason, FlagReason)
+        self.assertEqual(flag.reason.code, "user_in_segment")
+        self.assertIsInstance(flag.metadata, LegacyFlagMetadata)
+        self.assertIsNone(flag.metadata.payload)

posthog/test/test_feature_flags.py

@@ -38,7 +38,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "person-flag",
-                "is_simple_flag": True,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -69,6 +68,59 @@ class TestLocalEvaluation(unittest.TestCase):
         self.assertTrue(feature_flag_match)
         self.assertFalse(not_feature_flag_match)
 
+    def test_case_insensitive_matching(self):
+        self.client.feature_flags = [
+            {
+                "id": 1,
+                "name": "Beta Feature",
+                "key": "person-flag",
+                "is_simple_flag": True,
+                "active": True,
+                "filters": {
+                    "groups": [
+                        {
+                            "properties": [
+                                {
+                                    "key": "location",
+                                    "operator": "exact",
+                                    "value": ["Straße"],
+                                    "type": "person",
+                                }
+                            ],
+                            "rollout_percentage": 100,
+                        },
+                        {
+                            "properties": [
+                                {
+                                    "key": "star",
+                                    "operator": "exact",
+                                    "value": ["ſun"],
+                                    "type": "person",
+                                }
+                            ],
+                            "rollout_percentage": 100,
+                        },
+                    ],
+                },
+            }
+        ]
+
+        self.assertTrue(
+            self.client.get_feature_flag("person-flag", "some-distinct-id", person_properties={"location": "straße"})
+        )
+
+        self.assertTrue(
+            self.client.get_feature_flag("person-flag", "some-distinct-id", person_properties={"location": "strasse"})
+        )
+
+        self.assertTrue(
+            self.client.get_feature_flag("person-flag", "some-distinct-id", person_properties={"star": "ſun"})
+        )
+
+        self.assertTrue(
+            self.client.get_feature_flag("person-flag", "some-distinct-id", person_properties={"star": "sun"})
+        )
+
     @mock.patch("posthog.client.decide")
     @mock.patch("posthog.client.get")
     def test_flag_group_properties(self, patch_get, patch_decide):
@@ -77,7 +129,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "group-flag",
-                "is_simple_flag": True,
                 "active": True,
                 "filters": {
                     "aggregation_group_type_index": 0,
@@ -170,7 +221,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "complex-flag",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -286,7 +336,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": True,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -301,7 +350,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 2,
                 "name": "Beta Feature",
                 "key": "beta-feature2",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -343,7 +391,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": True,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -358,7 +405,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 2,
                 "name": "Beta Feature",
                 "key": "beta-feature2",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -410,7 +456,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": True,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -457,7 +502,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "rollout_percentage": 100,
                 "filters": {
@@ -472,7 +516,7 @@ class TestLocalEvaluation(unittest.TestCase):
             }
         ]
         # decide called always because experience_continuity is set
-        self.assertTrue(client.get_feature_flag("beta-feature", "distinct_id"), "decide-fallback-value")
+        self.assertEqual(client.get_feature_flag("beta-feature", "distinct_id"), "decide-fallback-value")
         self.assertEqual(patch_decide.call_count, 1)
 
     @mock.patch.object(Client, "capture")
@@ -487,7 +531,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "rollout_percentage": 100,
                 "filters": {
@@ -503,7 +546,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 2,
                 "name": "Beta Feature",
                 "key": "disabled-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -518,7 +560,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 3,
                 "name": "Beta Feature",
                 "key": "beta-feature2",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -551,7 +592,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "rollout_percentage": 100,
                 "filters": {
@@ -570,7 +610,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 2,
                 "name": "Beta Feature",
                 "key": "disabled-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -588,7 +627,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 3,
                 "name": "Beta Feature",
                 "key": "beta-feature2",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -654,7 +692,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "rollout_percentage": 100,
                 "filters": {
@@ -670,7 +707,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 2,
                 "name": "Beta Feature",
                 "key": "disabled-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -690,13 +726,11 @@ class TestLocalEvaluation(unittest.TestCase):
     @mock.patch.object(Client, "capture")
     @mock.patch("posthog.client.decide")
     def test_get_all_flags_and_payloads_with_no_fallback(self, patch_decide, patch_capture):
-        patch_decide.return_value = {"featureFlags": {"beta-feature": "variant-1", "beta-feature2": "variant-2"}}
         client = self.client
         basic_flag = {
             "id": 1,
             "name": "Beta Feature",
             "key": "beta-feature",
-            "is_simple_flag": False,
             "active": True,
             "rollout_percentage": 100,
             "filters": {
@@ -715,7 +749,6 @@ class TestLocalEvaluation(unittest.TestCase):
             "id": 2,
             "name": "Beta Feature",
             "key": "disabled-feature",
-            "is_simple_flag": False,
             "active": True,
             "filters": {
                 "groups": [
@@ -733,7 +766,6 @@ class TestLocalEvaluation(unittest.TestCase):
             basic_flag,
             disabled_flag,
         ]
-        client.feature_flags_by_key = {"beta-feature": basic_flag, "disabled-feature": disabled_flag}
         self.assertEqual(
             client.get_all_flags_and_payloads("distinct_id")["featureFlagPayloads"], {"beta-feature": "new"}
         )
@@ -751,7 +783,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "rollout_percentage": 100,
                 "filters": {
@@ -767,7 +798,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 2,
                 "name": "Beta Feature",
                 "key": "disabled-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -782,7 +812,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 3,
                 "name": "Beta Feature",
                 "key": "beta-feature2",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -814,7 +843,6 @@ class TestLocalEvaluation(unittest.TestCase):
             "id": 1,
             "name": "Beta Feature",
             "key": "beta-feature",
-            "is_simple_flag": False,
             "active": True,
             "rollout_percentage": 100,
             "filters": {
@@ -833,7 +861,6 @@ class TestLocalEvaluation(unittest.TestCase):
             "id": 2,
             "name": "Beta Feature",
             "key": "disabled-feature",
-            "is_simple_flag": False,
             "active": True,
             "filters": {
                 "groups": [
@@ -851,7 +878,6 @@ class TestLocalEvaluation(unittest.TestCase):
             "id": 3,
             "name": "Beta Feature",
             "key": "beta-feature2",
-            "is_simple_flag": False,
             "active": True,
             "filters": {
                 "groups": [
@@ -870,7 +896,6 @@ class TestLocalEvaluation(unittest.TestCase):
             flag_2,
             flag_3,
         ]
-        client.feature_flags_by_key = {"beta-feature": flag_1, "disabled-feature": flag_2, "beta-feature2": flag_3}
         # beta-feature2 has no value
         self.assertEqual(
             client.get_all_flags_and_payloads("distinct_id", only_evaluate_locally=True)["featureFlagPayloads"],
@@ -888,7 +913,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "rollout_percentage": 100,
                 "filters": {
@@ -904,7 +928,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 2,
                 "name": "Beta Feature",
                 "key": "disabled-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -927,7 +950,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": False,
                 "active": False,
                 "rollout_percentage": 100,
                 "filters": {
@@ -943,7 +965,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 2,
                 "name": "Beta Feature",
                 "key": "disabled-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -969,7 +990,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 id: 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": True,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -1007,13 +1027,15 @@ class TestLocalEvaluation(unittest.TestCase):
             "beta-feature",
             "some-distinct-id",
             person_properties={
-                "latestBuildVersion": "24.32..1",
+                "latestBuildVersion": "24.32.1",
                 "latestBuildVersionMajor": "24",
                 "latestBuildVersionMinor": "32",
                 "latestBuildVersionPatch": "1",
             },
         )
 
+        self.assertEqual(feature_flag_match, True)
+
     @mock.patch("posthog.client.decide")
     @mock.patch("posthog.client.get")
     def test_feature_flags_local_evaluation_for_cohorts(self, patch_get, patch_decide):
@@ -1023,7 +1045,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 2,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -1094,7 +1115,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 2,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -1207,7 +1227,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": True,
                 "active": True,
                 "rollout_percentage": 100,
                 "filters": {
@@ -1232,7 +1251,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": True,
                 "active": True,
                 "rollout_percentage": 0,
                 "filters": {
@@ -1257,7 +1275,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": True,
                 "active": True,
                 "rollout_percentage": None,
                 "filters": {
@@ -1281,7 +1298,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": True,
                 "active": True,
                 "rollout_percentage": 100,
                 "filters": {
@@ -1305,7 +1321,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "rollout_percentage": 100,
                 "filters": {
@@ -1330,7 +1345,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": True,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -1352,7 +1366,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "rollout_percentage": 100,
                 "filters": {
@@ -1418,7 +1431,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "rollout_percentage": 100,
                 "filters": {
@@ -1459,7 +1471,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "rollout_percentage": 100,
                 "filters": {
@@ -1511,7 +1522,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "rollout_percentage": 100,
                 "filters": {
@@ -1552,7 +1562,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "rollout_percentage": 100,
                 "filters": {
@@ -1595,7 +1604,6 @@ class TestLocalEvaluation(unittest.TestCase):
             "id": 1,
             "name": "Beta Feature",
             "key": "person-flag",
-            "is_simple_flag": True,
             "active": True,
             "filters": {
                 "groups": [
@@ -1615,7 +1623,6 @@ class TestLocalEvaluation(unittest.TestCase):
             },
         }
         self.client.feature_flags = [basic_flag]
-        self.client.feature_flags_by_key = {"person-flag": basic_flag}
 
         self.assertEqual(
             self.client.get_feature_flag_payload(
@@ -1659,7 +1666,6 @@ class TestLocalEvaluation(unittest.TestCase):
             "id": 1,
             "name": "Beta Feature",
             "key": "beta-feature",
-            "is_simple_flag": False,
             "active": True,
             "rollout_percentage": 100,
             "filters": {
@@ -1684,7 +1690,6 @@ class TestLocalEvaluation(unittest.TestCase):
             },
         }
         self.client.feature_flags = [multivariate_flag]
-        self.client.feature_flags_by_key = {"beta-feature": multivariate_flag}
 
         self.assertEqual(
             self.client.get_feature_flag_payload(
@@ -2237,7 +2242,6 @@ class TestCaptureCalls(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "complex-flag",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -2337,6 +2341,143 @@ class TestCaptureCalls(unittest.TestCase):
             disable_geoip=None,
         )
 
+    @mock.patch.object(Client, "capture")
+    @mock.patch("posthog.client.decide")
+    def test_capture_is_called_with_flag_details(self, patch_decide, patch_capture):
+        patch_decide.return_value = {
+            "flags": {
+                "decide-flag": {
+                    "key": "decide-flag",
+                    "enabled": True,
+                    "variant": "decide-variant",
+                    "reason": {
+                        "description": "Matched condition set 1",
+                    },
+                    "metadata": {
+                        "id": 23,
+                        "version": 42,
+                    },
+                }
+            },
+            "requestId": "18043bf7-9cf6-44cd-b959-9662ee20d371",
+        }
+        client = Client(FAKE_TEST_API_KEY)
+
+        self.assertEqual(client.get_feature_flag("decide-flag", "some-distinct-id"), "decide-variant")
+        self.assertEqual(patch_capture.call_count, 1)
+        patch_capture.assert_called_with(
+            "some-distinct-id",
+            "$feature_flag_called",
+            {
+                "$feature_flag": "decide-flag",
+                "$feature_flag_response": "decide-variant",
+                "locally_evaluated": False,
+                "$feature/decide-flag": "decide-variant",
+                "$feature_flag_reason": "Matched condition set 1",
+                "$feature_flag_id": 23,
+                "$feature_flag_version": 42,
+                "$feature_flag_request_id": "18043bf7-9cf6-44cd-b959-9662ee20d371",
+            },
+            groups={},
+            disable_geoip=None,
+        )
+
+    @mock.patch.object(Client, "capture")
+    @mock.patch("posthog.client.decide")
+    def test_capture_is_called_with_flag_details_and_payload(self, patch_decide, patch_capture):
+        patch_decide.return_value = {
+            "flags": {
+                "decide-flag-with-payload": {
+                    "key": "decide-flag-with-payload",
+                    "enabled": True,
+                    "variant": None,
+                    "reason": {
+                        "code": "matched_condition",
+                        "condition_index": 0,
+                        "description": "Matched condition set 1",
+                    },
+                    "metadata": {
+                        "id": 23,
+                        "version": 42,
+                        "payload": '{"foo": "bar"}',
+                    },
+                }
+            },
+            "requestId": "18043bf7-9cf6-44cd-b959-9662ee20d371",
+        }
+        client = Client(FAKE_TEST_API_KEY)
+
+        self.assertEqual(
+            client.get_feature_flag_payload("decide-flag-with-payload", "some-distinct-id"), '{"foo": "bar"}'
+        )
+        self.assertEqual(patch_capture.call_count, 1)
+        patch_capture.assert_called_with(
+            "some-distinct-id",
+            "$feature_flag_called",
+            {
+                "$feature_flag": "decide-flag-with-payload",
+                "$feature_flag_response": True,
+                "locally_evaluated": False,
+                "$feature/decide-flag-with-payload": True,
+                "$feature_flag_reason": "Matched condition set 1",
+                "$feature_flag_id": 23,
+                "$feature_flag_version": 42,
+                "$feature_flag_request_id": "18043bf7-9cf6-44cd-b959-9662ee20d371",
+                "$feature_flag_payload": '{"foo": "bar"}',
+            },
+            groups={},
+            disable_geoip=None,
+        )
+
+    @mock.patch("posthog.client.decide")
+    def test_capture_is_called_but_does_not_add_all_flags(self, patch_decide):
+        patch_decide.return_value = {"featureFlags": {"decide-flag": "decide-value"}}
+        client = Client(FAKE_TEST_API_KEY, personal_api_key=FAKE_TEST_API_KEY)
+        client.feature_flags = [
+            {
+                "id": 1,
+                "name": "Beta Feature",
+                "key": "complex-flag",
+                "active": True,
+                "filters": {
+                    "groups": [
+                        {
+                            "properties": [{"key": "region", "value": "USA"}],
+                            "rollout_percentage": 100,
+                        },
+                    ],
+                },
+            },
+            {
+                "id": 2,
+                "name": "Gamma Feature",
+                "key": "simple-flag",
+                "active": True,
+                "filters": {
+                    "groups": [
+                        {
+                            "properties": [],
+                            "rollout_percentage": 100,
+                        },
+                    ],
+                },
+            },
+        ]
+
+        self.assertTrue(
+            client.get_feature_flag("complex-flag", "some-distinct-id", person_properties={"region": "USA"})
+        )
+
+        # Grab the capture message that was just added to the queue
+        msg = client.queue.get(block=False)
+        assert msg["event"] == "$feature_flag_called"
+        assert msg["properties"]["$feature_flag"] == "complex-flag"
+        assert msg["properties"]["$feature_flag_response"] is True
+        assert msg["properties"]["locally_evaluated"] is True
+        assert msg["properties"]["$feature/complex-flag"] is True
+        assert "$feature/simple-flag" not in msg["properties"]
+        assert "$active_feature_flags" not in msg["properties"]
+
     @mock.patch.object(Client, "capture")
     @mock.patch("posthog.client.decide")
     def test_capture_is_called_in_get_feature_flag_payload(self, patch_decide, patch_capture):
@@ -2351,7 +2492,6 @@ class TestCaptureCalls(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "person-flag",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -2429,7 +2569,6 @@ class TestCaptureCalls(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "complex-flag",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -2472,7 +2611,6 @@ class TestCaptureCalls(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "complex-flag",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -2531,7 +2669,6 @@ class TestConsistency(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "simple-flag",
-                "is_simple_flag": True,
                 "active": True,
                 "filters": {
                     "groups": [{"properties": [], "rollout_percentage": 45}],
@@ -3559,7 +3696,6 @@ class TestConsistency(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "multivariate-flag",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [{"properties": [], "rollout_percentage": 55}],
@@ -4587,3 +4723,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"))

posthog/test/test_request.py

@@ -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",

posthog/test/test_types.py

@@ -0,0 +1,175 @@
+import unittest
+
+from parameterized import parameterized
+
+from posthog.types import (
+    FeatureFlag,
+    FlagMetadata,
+    FlagReason,
+    LegacyFlagMetadata,
+    normalize_decide_response,
+    to_flags_and_payloads,
+)
+
+
+class TestTypes(unittest.TestCase):
+    @parameterized.expand([(True,), (False,)])
+    def test_normalize_decide_response_v4(self, has_errors: bool):
+        resp = {
+            "flags": {
+                "my-flag": FeatureFlag(
+                    key="my-flag",
+                    enabled=True,
+                    variant="test-variant",
+                    reason=FlagReason(
+                        code="matched_condition", condition_index=0, description="Matched condition set 1"
+                    ),
+                    metadata=FlagMetadata(id=1, payload='{"some": "json"}', version=2, description="test-description"),
+                )
+            },
+            "errorsWhileComputingFlags": has_errors,
+            "requestId": "test-id",
+        }
+
+        result = normalize_decide_response(resp)
+
+        flag = result["flags"]["my-flag"]
+        self.assertEqual(flag.key, "my-flag")
+        self.assertTrue(flag.enabled)
+        self.assertEqual(flag.variant, "test-variant")
+        self.assertEqual(flag.get_value(), "test-variant")
+        self.assertEqual(
+            flag.reason, FlagReason(code="matched_condition", condition_index=0, description="Matched condition set 1")
+        )
+        self.assertEqual(
+            flag.metadata, FlagMetadata(id=1, payload='{"some": "json"}', version=2, description="test-description")
+        )
+        self.assertEqual(result["errorsWhileComputingFlags"], has_errors)
+        self.assertEqual(result["requestId"], "test-id")
+
+    def test_normalize_decide_response_legacy(self):
+        # Test legacy response format with "featureFlags" and "featureFlagPayloads"
+        resp = {
+            "featureFlags": {"my-flag": "test-variant"},
+            "featureFlagPayloads": {"my-flag": '{"some": "json-payload"}'},
+            "errorsWhileComputingFlags": False,
+            "requestId": "test-id",
+        }
+
+        result = normalize_decide_response(resp)
+
+        flag = result["flags"]["my-flag"]
+        self.assertEqual(flag.key, "my-flag")
+        self.assertTrue(flag.enabled)
+        self.assertEqual(flag.variant, "test-variant")
+        self.assertEqual(flag.get_value(), "test-variant")
+        self.assertIsNone(flag.reason)
+        self.assertEqual(flag.metadata, LegacyFlagMetadata(payload='{"some": "json-payload"}'))
+        self.assertFalse(result["errorsWhileComputingFlags"])
+        self.assertEqual(result["requestId"], "test-id")
+        # Verify legacy fields are removed
+        self.assertNotIn("featureFlags", result)
+        self.assertNotIn("featureFlagPayloads", result)
+
+    def test_normalize_decide_response_boolean_flag(self):
+        # Test legacy response with boolean flag
+        resp = {"featureFlags": {"my-flag": True}, "errorsWhileComputingFlags": False}
+
+        result = normalize_decide_response(resp)
+
+        self.assertIn("requestId", result)
+        self.assertIsNone(result["requestId"])
+
+        flag = result["flags"]["my-flag"]
+        self.assertEqual(flag.key, "my-flag")
+        self.assertTrue(flag.enabled)
+        self.assertIsNone(flag.variant)
+        self.assertIsNone(flag.reason)
+        self.assertEqual(flag.metadata, LegacyFlagMetadata(payload=None))
+        self.assertFalse(result["errorsWhileComputingFlags"])
+        self.assertNotIn("featureFlags", result)
+        self.assertNotIn("featureFlagPayloads", result)
+
+    def test_to_flags_and_payloads_v4(self):
+        # Test v4 response format
+        resp = {
+            "flags": {
+                "my-variant-flag": FeatureFlag(
+                    key="my-variant-flag",
+                    enabled=True,
+                    variant="test-variant",
+                    reason=FlagReason(
+                        code="matched_condition", condition_index=0, description="Matched condition set 1"
+                    ),
+                    metadata=FlagMetadata(id=1, payload='{"some": "json"}', version=2, description="test-description"),
+                ),
+                "my-boolean-flag": FeatureFlag(
+                    key="my-boolean-flag",
+                    enabled=True,
+                    variant=None,
+                    reason=FlagReason(
+                        code="matched_condition", condition_index=0, description="Matched condition set 1"
+                    ),
+                    metadata=FlagMetadata(id=1, payload=None, version=2, description="test-description"),
+                ),
+                "disabled-flag": FeatureFlag(
+                    key="disabled-flag",
+                    enabled=False,
+                    variant=None,
+                    reason=None,
+                    metadata=LegacyFlagMetadata(payload=None),
+                ),
+            },
+            "errorsWhileComputingFlags": False,
+            "requestId": "test-id",
+        }
+
+        result = to_flags_and_payloads(resp)
+
+        self.assertEqual(result["featureFlags"]["my-variant-flag"], "test-variant")
+        self.assertEqual(result["featureFlags"]["my-boolean-flag"], True)
+        self.assertEqual(result["featureFlags"]["disabled-flag"], False)
+        self.assertEqual(result["featureFlagPayloads"]["my-variant-flag"], '{"some": "json"}')
+        self.assertNotIn("my-boolean-flag", result["featureFlagPayloads"])
+        self.assertNotIn("disabled-flag", result["featureFlagPayloads"])
+
+    def test_to_flags_and_payloads_empty(self):
+        # Test empty response
+        resp = {
+            "flags": {},
+            "errorsWhileComputingFlags": False,
+            "requestId": "test-id",
+        }
+
+        result = to_flags_and_payloads(resp)
+
+        self.assertEqual(result["featureFlags"], {})
+        self.assertEqual(result["featureFlagPayloads"], {})
+
+    def test_to_flags_and_payloads_with_payload(self):
+        resp = {
+            "flags": {
+                "decide-flag": {
+                    "key": "decide-flag",
+                    "enabled": True,
+                    "variant": "decide-variant",
+                    "reason": {
+                        "code": "matched_condition",
+                        "condition_index": 0,
+                        "description": "Matched condition set 1",
+                    },
+                    "metadata": {
+                        "id": 23,
+                        "version": 42,
+                        "payload": '{"foo": "bar"}',
+                    },
+                }
+            },
+            "requestId": "18043bf7-9cf6-44cd-b959-9662ee20d371",
+        }
+
+        normalized = normalize_decide_response(resp)
+        result = to_flags_and_payloads(normalized)
+
+        self.assertEqual(result["featureFlags"]["decide-flag"], "decide-variant")
+        self.assertEqual(result["featureFlagPayloads"]["decide-flag"], '{"foo": "bar"}')

posthog/test/test_utils.py

@@ -1,10 +1,14 @@
 import unittest
+from dataclasses import dataclass
 from datetime import date, datetime, timedelta
 from decimal import Decimal
+from typing import Optional
 from uuid import UUID
 
 import six
 from dateutil.tz import tzutc
+from pydantic import BaseModel
+from pydantic.v1 import BaseModel as BaseModelV1
 
 from posthog import utils
 
@@ -53,7 +57,10 @@ class TestUtils(unittest.TestCase):
         self.assertEqual(combined.keys(), pre_clean_keys)
 
         # test UUID separately, as the UUID object doesn't equal its string representation according to Python
-        self.assertEqual(utils.clean(UUID("12345678123456781234567812345678")), "12345678-1234-5678-1234-567812345678")
+        self.assertEqual(
+            utils.clean(UUID("12345678123456781234567812345678")),
+            "12345678-1234-5678-1234-567812345678",
+        )
 
     def test_clean_with_dates(self):
         dict_with_dates = {
@@ -81,6 +88,74 @@ class TestUtils(unittest.TestCase):
         self.assertEqual("http://posthog.io", utils.remove_trailing_slash("http://posthog.io/"))
         self.assertEqual("http://posthog.io", utils.remove_trailing_slash("http://posthog.io"))
 
+    def test_clean_pydantic(self):
+        class ModelV2(BaseModel):
+            foo: str
+            bar: int
+            baz: Optional[str] = None
+
+        class ModelV1(BaseModelV1):
+            foo: int
+            bar: str
+
+        class NestedModel(BaseModel):
+            foo: ModelV2
+
+        self.assertEqual(utils.clean(ModelV2(foo="1", bar=2)), {"foo": "1", "bar": 2, "baz": None})
+        self.assertEqual(utils.clean(ModelV1(foo=1, bar="2")), {"foo": 1, "bar": "2"})
+        self.assertEqual(
+            utils.clean(NestedModel(foo=ModelV2(foo="1", bar=2, baz="3"))),
+            {"foo": {"foo": "1", "bar": 2, "baz": "3"}},
+        )
+
+        class Dummy:
+            def model_dump(self, required_param):
+                pass
+
+        # Skips a class with a defined non-Pydantic `model_dump` method.
+        self.assertEqual(utils.clean({"test": Dummy()}), {})
+
+    def test_clean_dataclass(self):
+        @dataclass
+        class InnerDataClass:
+            inner_foo: str
+            inner_bar: int
+            inner_uuid: UUID
+            inner_date: datetime
+            inner_optional: Optional[str] = None
+
+        @dataclass
+        class TestDataClass:
+            foo: str
+            bar: int
+            nested: InnerDataClass
+
+        self.assertEqual(
+            utils.clean(
+                TestDataClass(
+                    foo="1",
+                    bar=2,
+                    nested=InnerDataClass(
+                        inner_foo="3",
+                        inner_bar=4,
+                        inner_uuid=UUID("12345678123456781234567812345678"),
+                        inner_date=datetime(2025, 1, 1),
+                    ),
+                )
+            ),
+            {
+                "foo": "1",
+                "bar": 2,
+                "nested": {
+                    "inner_foo": "3",
+                    "inner_bar": 4,
+                    "inner_uuid": "12345678-1234-5678-1234-567812345678",
+                    "inner_date": datetime(2025, 1, 1),
+                    "inner_optional": None,
+                },
+            },
+        )
+
 
 class TestSizeLimitedDict(unittest.TestCase):
     def test_size_limited_dict(self):

posthog/types.py

@@ -0,0 +1,171 @@
+from dataclasses import dataclass
+from typing import Any, List, Optional, TypedDict, Union, cast
+
+FlagValue = Union[bool, str]
+
+
+@dataclass(frozen=True)
+class FlagReason:
+    code: str
+    condition_index: int
+    description: str
+
+    @classmethod
+    def from_json(cls, resp: Any) -> Optional["FlagReason"]:
+        if not resp:
+            return None
+        return cls(
+            code=resp.get("code", ""),
+            condition_index=resp.get("condition_index", 0),
+            description=resp.get("description", ""),
+        )
+
+
+@dataclass(frozen=True)
+class LegacyFlagMetadata:
+    payload: Any
+
+
+@dataclass(frozen=True)
+class FlagMetadata:
+    id: int
+    payload: Optional[str]
+    version: int
+    description: str
+
+    @classmethod
+    def from_json(cls, resp: Any) -> Union["FlagMetadata", LegacyFlagMetadata]:
+        if not resp:
+            return LegacyFlagMetadata(payload=None)
+        return cls(
+            id=resp.get("id", 0),
+            payload=resp.get("payload"),
+            version=resp.get("version", 0),
+            description=resp.get("description", ""),
+        )
+
+
+@dataclass(frozen=True)
+class FeatureFlag:
+    key: str
+    enabled: bool
+    variant: Optional[str]
+    reason: Optional[FlagReason]
+    metadata: Union[FlagMetadata, LegacyFlagMetadata]
+
+    def get_value(self) -> FlagValue:
+        return self.variant or self.enabled
+
+    @classmethod
+    def from_json(cls, resp: Any) -> "FeatureFlag":
+        reason = None
+        if resp.get("reason"):
+            reason = FlagReason.from_json(resp.get("reason"))
+
+        metadata = None
+        if resp.get("metadata"):
+            metadata = FlagMetadata.from_json(resp.get("metadata"))
+        else:
+            metadata = LegacyFlagMetadata(payload=None)
+
+        return cls(
+            key=resp.get("key"),
+            enabled=resp.get("enabled"),
+            variant=resp.get("variant"),
+            reason=reason,
+            metadata=metadata,
+        )
+
+    @classmethod
+    def from_value_and_payload(cls, key: str, value: FlagValue, payload: Any) -> "FeatureFlag":
+        enabled, variant = (True, value) if isinstance(value, str) else (value, None)
+        return cls(
+            key=key,
+            enabled=enabled,
+            variant=variant,
+            reason=None,
+            metadata=LegacyFlagMetadata(
+                payload=payload if payload else None,
+            ),
+        )
+
+
+class DecideResponse(TypedDict, total=False):
+    flags: dict[str, FeatureFlag]
+    errorsWhileComputingFlags: bool
+    requestId: str
+    quotaLimit: Optional[List[str]]
+
+
+class FlagsAndPayloads(TypedDict, total=True):
+    featureFlags: Optional[dict[str, FlagValue]]
+    featureFlagPayloads: Optional[dict[str, Any]]
+
+
+def normalize_decide_response(resp: Any) -> DecideResponse:
+    """
+    Normalize the response from the decide API endpoint into a v4 DecideResponse.
+
+    Args:
+        resp: A v3 or v4 response from the decide API endpoint.
+
+    Returns:
+        A DecideResponse containing feature flags and their details.
+    """
+    if "requestId" not in resp:
+        resp["requestId"] = None
+    if "flags" in resp:
+        flags = resp["flags"]
+        # For each flag, create a FeatureFlag object
+        for key, value in flags.items():
+            if isinstance(value, FeatureFlag):
+                continue
+            value["key"] = key
+            flags[key] = FeatureFlag.from_json(value)
+    else:
+        # Handle legacy format
+        featureFlags = resp.get("featureFlags", {})
+        featureFlagPayloads = resp.get("featureFlagPayloads", {})
+        resp.pop("featureFlags", None)
+        resp.pop("featureFlagPayloads", None)
+        # look at each key in featureFlags and create a FeatureFlag object
+        flags = {}
+        for key, value in featureFlags.items():
+            flags[key] = FeatureFlag.from_value_and_payload(key, value, featureFlagPayloads.get(key, None))
+        resp["flags"] = flags
+    return cast(DecideResponse, resp)
+
+
+def to_flags_and_payloads(resp: DecideResponse) -> FlagsAndPayloads:
+    """
+    Convert a DecideResponse into a FlagsAndPayloads object which is a
+    dict of feature flags and their payloads. This is needed by certain
+    functions in the client.
+    Args:
+        resp: A DecideResponse containing feature flags and their payloads.
+
+    Returns:
+        A tuple containing:
+            - A dictionary mapping flag keys to their values (bool or str)
+            - A dictionary mapping flag keys to their payloads
+    """
+    return {"featureFlags": to_values(resp), "featureFlagPayloads": to_payloads(resp)}
+
+
+def to_values(response: DecideResponse) -> Optional[dict[str, FlagValue]]:
+    if "flags" not in response:
+        return None
+
+    flags = response.get("flags", {})
+    return {key: value.get_value() for key, value in flags.items() if isinstance(value, FeatureFlag)}
+
+
+def to_payloads(response: DecideResponse) -> Optional[dict[str, str]]:
+    if "flags" not in response:
+        return None
+
+    return {
+        key: value.metadata.payload
+        for key, value in response.get("flags", {}).items()
+        if isinstance(value, FeatureFlag) and value.enabled and value.metadata.payload
+    }

posthog/utils.py

@@ -2,6 +2,7 @@ import logging
 import numbers
 import re
 from collections import defaultdict
+from dataclasses import asdict, is_dataclass
 from datetime import date, datetime, timezone
 from decimal import Decimal
 from uuid import UUID
@@ -51,14 +52,26 @@ def clean(item):
         return float(item)
     if isinstance(item, UUID):
         return str(item)
-    elif isinstance(item, (six.string_types, bool, numbers.Number, datetime, date, type(None))):
+    if isinstance(item, (six.string_types, bool, numbers.Number, datetime, date, type(None))):
         return item
-    elif isinstance(item, (set, list, tuple)):
+    if isinstance(item, (set, list, tuple)):
         return _clean_list(item)
-    elif isinstance(item, dict):
+    # Pydantic model
+    try:
+        # v2+
+        if hasattr(item, "model_dump") and callable(item.model_dump):
+            item = item.model_dump()
+        # v1
+        elif hasattr(item, "dict") and callable(item.dict):
+            item = item.dict()
+    except TypeError as e:
+        log.debug(f"Could not serialize Pydantic-like model: {e}")
+        pass
+    if isinstance(item, dict):
         return _clean_dict(item)
-    else:
-        return _coerce_unicode(item)
+    if is_dataclass(item) and not isinstance(item, type):
+        return _clean_dataclass(item)
+    return _coerce_unicode(item)
 
 
 def _clean_list(list_):
@@ -80,6 +93,12 @@ def _clean_dict(dict_):
     return data
 
 
+def _clean_dataclass(dataclass_):
+    data = asdict(dataclass_)
+    data = _clean_dict(data)
+    return data
+
+
 def _coerce_unicode(cmplx):
     try:
         item = cmplx.decode("utf-8", "strict")
@@ -115,3 +134,43 @@ def convert_to_datetime_aware(date_obj):
     if date_obj.tzinfo is None:
         date_obj = date_obj.replace(tzinfo=timezone.utc)
     return date_obj
+
+
+def str_icontains(source, search):
+    """
+    Check if a string contains another string, ignoring case.
+
+    Args:
+        source: The string to search within
+        search: The substring to search for
+
+    Returns:
+        bool: True if search is a substring of source (case-insensitive), False otherwise
+
+    Examples:
+        >>> str_icontains("Hello World", "WORLD")
+        True
+        >>> str_icontains("Hello World", "python")
+        False
+    """
+    return str(search).casefold() in str(source).casefold()
+
+
+def str_iequals(value, comparand):
+    """
+    Check if a string equals another string, ignoring case.
+
+    Args:
+        value: The string to compare
+        comparand: The string to compare with
+
+    Returns:
+        bool: True if value and comparand are equal (case-insensitive), False otherwise
+
+    Examples:
+        >>> str_iequals("Hello World", "hello world")
+        True
+        >>> str_iequals("Hello World", "hello")
+        False
+    """
+    return str(value).casefold() == str(comparand).casefold()

posthog/version.py

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

setup.cfg

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

setup.py

@@ -14,18 +14,54 @@ 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+
+]
 
 extras_require = {
     "dev": [
         "black",
+        "django-stubs",
         "isort",
         "flake8",
         "flake8-print",
+        "lxml",
+        "mypy",
+        "mypy-baseline",
+        "types-mock",
+        "types-python-dateutil",
+        "types-requests",
+        "types-setuptools",
+        "types-six",
         "pre-commit",
+        "pydantic",
+    ],
+    "test": [
+        "mock>=2.0.0",
+        "freezegun==1.5.1",
+        "pylint",
+        "flake8",
+        "coverage",
+        "pytest",
+        "pytest-timeout",
+        "pytest-asyncio",
+        "django",
+        "openai",
+        "anthropic",
+        "langgraph",
+        "langchain-community>=0.2.0",
+        "langchain-openai>=0.2.0",
+        "langchain-anthropic>=0.2.0",
+        "pydantic",
+        "parameterized>=0.8.1",
     ],
-    "test": ["mock>=2.0.0", "freezegun==0.3.15", "pylint", "flake8", "coverage", "pytest", "pytest-timeout", "django"],
     "sentry": ["sentry-sdk", "django"],
+    "langchain": ["langchain>=0.2.0"],
 }
 
 setup(
@@ -37,7 +73,16 @@ setup(
     maintainer="PostHog",
     maintainer_email="hey@posthog.com",
     test_suite="posthog.test.all",
-    packages=["posthog", "posthog.test", "posthog.sentry", "posthog.exception_integrations"],
+    packages=[
+        "posthog",
+        "posthog.ai",
+        "posthog.ai.langchain",
+        "posthog.ai.openai",
+        "posthog.ai.anthropic",
+        "posthog.test",
+        "posthog.sentry",
+        "posthog.exception_integrations",
+    ],
     license="MIT License",
     install_requires=install_requires,
     extras_require=extras_require,
@@ -49,16 +94,10 @@ setup(
         "License :: OSI Approved :: MIT License",
         "Operating System :: OS Independent",
         "Programming Language :: Python",
-        "Programming Language :: Python :: 2",
-        "Programming Language :: Python :: 2.6",
-        "Programming Language :: Python :: 2.7",
-        "Programming Language :: Python :: 3",
-        "Programming Language :: Python :: 3.2",
-        "Programming Language :: Python :: 3.3",
-        "Programming Language :: Python :: 3.4",
-        "Programming Language :: Python :: 3.5",
-        "Programming Language :: Python :: 3.6",
-        "Programming Language :: Python :: 3.7",
-        "Programming Language :: Python :: 3.8",
+        "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
+        "Programming Language :: Python :: 3.12",
+        "Programming Language :: Python :: 3.13",
     ],
 )

setup_analytics.py

@@ -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"]
 
@@ -29,6 +36,10 @@ setup(
     test_suite="posthoganalytics.test.all",
     packages=[
         "posthoganalytics",
+        "posthoganalytics.ai",
+        "posthoganalytics.ai.langchain",
+        "posthoganalytics.ai.openai",
+        "posthoganalytics.ai.anthropic",
         "posthoganalytics.test",
         "posthoganalytics.sentry",
         "posthoganalytics.exception_integrations",
@@ -58,5 +69,8 @@ setup(
         "Programming Language :: Python :: 3.6",
         "Programming Language :: Python :: 3.7",
         "Programming Language :: Python :: 3.8",
+        "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
     ],
 )

simulator.py

@@ -24,7 +24,6 @@ parser.add_argument("--type", help="The posthog message type")
 
 parser.add_argument("--distinct_id", help="the user id to send the event as")
 parser.add_argument("--anonymousId", help="the anonymous user id to send the event as")
-parser.add_argument("--context", help="additional context for the event (JSON-encoded)")
 
 parser.add_argument("--event", help="the event name to send with the event")
 parser.add_argument("--properties", help="the event properties to send (JSON-encoded)")
@@ -48,7 +47,6 @@ def capture():
         options.event,
         anonymous_id=options.anonymousId,
         properties=json_hash(options.properties),
-        context=json_hash(options.context),
     )
 
 
@@ -58,7 +56,6 @@ def page():
         name=options.name,
         anonymous_id=options.anonymousId,
         properties=json_hash(options.properties),
-        context=json_hash(options.context),
     )
 
 
@@ -67,7 +64,6 @@ def identify():
         options.distinct_id,
         anonymous_id=options.anonymousId,
         traits=json_hash(options.traits),
-        context=json_hash(options.context),
     )
 
 
@@ -75,7 +71,6 @@ def set_once():
     posthog.set_once(
         options.distinct_id,
         properties=json_hash(options.traits),
-        context=json_hash(options.context),
     )
 
 
@@ -83,7 +78,6 @@ def set():
     posthog.set(
         options.distinct_id,
         properties=json_hash(options.traits),
-        context=json_hash(options.context),
     )