changelog

* maybe

* bump version

.github/workflows/ci.yml

@@ -9,19 +9,19 @@ jobs:
         runs-on: ubuntu-latest
 
         steps:
-            - uses: actions/checkout@v2
+            - uses: actions/checkout@85e6279cec87321a52edac9c87bce653a07cf6c2
               with:
                   fetch-depth: 1
 
             - name: Set up Python 3.11
-              uses: actions/setup-python@v2
+              uses: actions/setup-python@8d9ed9ac5c53483de85588cdf95a591a75ab9f55
               with:
                   python-version: 3.11.11
 
-            - uses: actions/cache@v3
+            - uses: actions/cache@5a3ec84eff668545956fd18022155c47e93e2684
               with:
                   path: ~/.cache/pip
-                  key: ${{ runner.os }}-pip-${{ hashFiles('setup.py') }}
+                  key: ${{ runner.os }}-pip-${{ hashFiles('pyproject.toml') }}
                   restore-keys: |
                       ${{ runner.os }}-pip-
 
@@ -30,17 +30,9 @@ jobs:
                   python -m pip install -e .[dev]
               if: steps.cache.outputs.cache-hit != 'true'
 
-            - name: Check formatting with black
+            - name: Check formatting with ruff
               run: |
-                  black --check .
-
-            - name: Lint with flake8
-              run: |
-                  flake8 posthog --ignore E501,W503
-
-            - name: Check import order with isort
-              run: |
-                  isort --check-only .
+                  ruff format --check .
 
             - name: Check types with mypy
               run: |
@@ -54,12 +46,12 @@ jobs:
                 python-version: ['3.9', '3.10', '3.11', '3.12', '3.13']
 
         steps:
-            - uses: actions/checkout@v2
+            - uses: actions/checkout@85e6279cec87321a52edac9c87bce653a07cf6c2
               with:
                   fetch-depth: 1
 
             - name: Set up Python ${{ matrix.python-version }}
-              uses: actions/setup-python@v2
+              uses: actions/setup-python@8d9ed9ac5c53483de85588cdf95a591a75ab9f55
               with:
                   python-version: ${{ matrix.python-version }}
 

.github/workflows/release.yaml

@@ -17,25 +17,25 @@ jobs:
       TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }}
     steps:
       - name: Checkout the repository
-        uses: actions/checkout@v2
+        uses: actions/checkout@85e6279cec87321a52edac9c87bce653a07cf6c2
         with:
           fetch-depth: 0
           token: ${{ secrets.POSTHOG_BOT_GITHUB_TOKEN }}
 
       - name: Set up Python
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@8d9ed9ac5c53483de85588cdf95a591a75ab9f55
 
       - name: Detect version
         run: echo "REPO_VERSION=$(python3 posthog/version.py)" >> $GITHUB_ENV
 
       - name: Prepare for building release
-        run: pip install -U pip setuptools wheel twine
+        run: pip install -U pip setuptools packaging wheel twine
 
       - name: Push release to PyPI
         run: make release && make release_analytics
 
       - name: Create GitHub release
-        uses: actions/create-release@v1
+        uses: actions/create-release@0cb9c9b65d5d1901c1f53e5e66eaf4afd303e70e # v1
         env:
           GITHUB_TOKEN: ${{ secrets.POSTHOG_BOT_GITHUB_TOKEN }}
         with:

.pre-commit-config.yaml

@@ -1,9 +1,10 @@
 repos:
-  - repo: https://github.com/psf/black
-    rev: stable
-    hooks:
-    - id: black
-  - repo: https://github.com/pycqa/isort
-    rev: 5.7.0
-    hooks:
-      - id: isort
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  # Ruff version.
+  rev: v0.11.12
+  hooks:
+    # Run the linter.
+    - id: ruff-check
+      args: [ --fix ]
+    # Run the formatter.
+    - id: ruff-format

CHANGELOG.md

@@ -1,3 +1,32 @@
+## 4.3.3 - 2025-06-06
+
+Add `setup()` function to initialise default client
+
+## 4.3.2 - 2025-06-06
+
+Add context management:
+
+- New context manager with `posthog.new_context()`
+- Tag functions: `posthog.tag()`, `posthog.get_tags()`, `posthog.clear_tags()`
+- Function decorator:
+  - `@posthog.scoped` - Creates context and captures exceptions thrown within the function
+- Automatic deduplication of exceptions to ensure each exception is only captured once
+
+## 4.2.1 - 2025-6-05
+
+1. fix: feature flag request use geoip_disable (#235)
+2. chore: pin actions versions (#210)
+3. fix: opinionated setup and clean fn fix (#240)
+4. fix: release action failed (#241)
+
+## 4.2.0 - 2025-05-22
+
+Add support for google gemini
+
+## 4.1.0 - 2025-05-22
+
+Moved ai openai package to a composition approach over inheritance.
+
 ## 4.0.1 – 2025-04-29
 
 1. Remove deprecated `monotonic` library. Use Python's core `time.monotonic` function instead

Makefile

@@ -1,6 +1,5 @@
 lint:
-	pylint --rcfile=.pylintrc --reports=y --exit-zero analytics | tee pylint.out
-	flake8 --max-complexity=10 --statistics analytics > flake8.out || true
+	uvx ruff format
 
 test:
 	coverage run -m pytest

README.md

@@ -1,7 +1,14 @@
 # PostHog Python
 
-[![PyPI](https://img.shields.io/pypi/v/posthog)](https://pypi.org/project/posthog/)
-
+<p align="center">
+  <img alt="posthoglogo" src="https://user-images.githubusercontent.com/65415371/205059737-c8a4f836-4889-4654-902e-f302b187b6a0.png">
+</p>
+<p align="center">
+   <a href="https://pypi.org/project/posthog/"><img alt="pypi installs" src="https://img.shields.io/pypi/v/posthog"/></a>
+   <img alt="GitHub contributors" src="https://img.shields.io/github/contributors/posthog/posthog-python">
+  <img alt="GitHub commit activity" src="https://img.shields.io/github/commit-activity/m/posthog/posthog-python"/>
+  <img alt="GitHub closed issues" src="https://img.shields.io/github/issues-closed/posthog/posthog-python"/>
+</p>
 
 Please see the [Python integration docs](https://posthog.com/docs/integrations/python-integration) for details.
 
@@ -14,9 +21,22 @@ Please see the [Python integration docs](https://posthog.com/docs/integrations/p
 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`
+4. you have to run `pre-commit install` to have auto linting pre commit
+5. Run `make test`
   1. To run a specific test do `pytest -k test_no_api_key`
 
+## PostHog recommends `uv` so...
+
+```bash
+uv python install 3.9.19
+uv python pin 3.9.19
+uv venv env
+source env/bin/activate
+uv pip install --editable ".[dev,test]"
+pre-commit install
+make test
+```
+
 ### Running Locally
 
 Assuming you have a [local version of PostHog](https://posthog.com/docs/developing-locally) running, you can run `python3 example.py` to see the library in action.
@@ -41,8 +61,4 @@ Then navigate to `http://127.0.0.1:8080/sentry-debug/` and you should get an eve
 
 ### Releasing Versions
 
-Updated are released using GitHub Actions: after bumping `version.py` in `master` and adding to `CHANGELOG.md`, go to [our release workflow's page](https://github.com/PostHog/posthog-python/actions/workflows/release.yaml) and dispatch it manually, using workflow from `master`.
-
-## Questions?
-
-### [Join our Slack community.](https://join.slack.com/t/posthogusers/shared_invite/enQtOTY0MzU5NjAwMDY3LTc2MWQ0OTZlNjhkODk3ZDI3NDVjMDE1YjgxY2I4ZjI4MzJhZmVmNjJkN2NmMGJmMzc2N2U3Yjc3ZjI5NGFlZDQ)
+Updated are released using GitHub Actions: after bumping `version.py` in `master` and adding to `CHANGELOG.md`, go to [our release workflow's page](https://github.com/PostHog/posthog-python/actions/workflows/release.yaml) and dispatch it manually, using workflow from `master`.

example.py

@@ -6,7 +6,9 @@ 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)"
+    "--flag",
+    default="person-on-events-enabled",
+    help="Feature flag key to check (default: person-on-events-enabled)",
 )
 args = parser.parse_args()
 
@@ -38,10 +40,19 @@ print(
 
 
 # Capture an event
-posthog.capture("distinct_id", "event", {"property1": "value", "property2": "value"}, send_feature_flags=True)
+posthog.capture(
+    "distinct_id",
+    "event",
+    {"property1": "value", "property2": "value"},
+    send_feature_flags=True,
+)
 
 print(posthog.feature_enabled("beta-feature", "distinct_id"))
-print(posthog.feature_enabled("beta-feature-groups", "distinct_id", groups={"company": "id:5"}))
+print(
+    posthog.feature_enabled(
+        "beta-feature-groups", "distinct_id", groups={"company": "id:5"}
+    )
+)
 
 print(posthog.feature_enabled("beta-feature", "distinct_id"))
 
@@ -53,9 +64,14 @@ exit()
 
 posthog.alias("distinct_id", "new_distinct_id")
 
-posthog.capture("new_distinct_id", "event2", {"property1": "value", "property2": "value"})
 posthog.capture(
-    "new_distinct_id", "event-with-groups", {"property1": "value", "property2": "value"}, groups={"company": "id:5"}
+    "new_distinct_id", "event2", {"property1": "value", "property2": "value"}
+)
+posthog.capture(
+    "new_distinct_id",
+    "event-with-groups",
+    {"property1": "value", "property2": "value"},
+    groups={"company": "id:5"},
 )
 
 # # Add properties to the person
@@ -82,7 +98,13 @@ posthog.set("new_distinct_id", {"current_browser": "Firefox"})
 # Local Evaluation
 
 # If flag has City=Sydney, this call doesn't go to `/decide`
-print(posthog.feature_enabled("test-flag", "distinct_id_random_22", person_properties={"$geoip_city_name": "Sydney"}))
+print(
+    posthog.feature_enabled(
+        "test-flag",
+        "distinct_id_random_22",
+        person_properties={"$geoip_city_name": "Sydney"},
+    )
+)
 
 print(
     posthog.feature_enabled(
@@ -98,10 +120,52 @@ print(posthog.get_all_flags("distinct_id_random_22"))
 print(posthog.get_all_flags("distinct_id_random_22", only_evaluate_locally=True))
 print(
     posthog.get_all_flags(
-        "distinct_id_random_22", person_properties={"$geoip_city_name": "Sydney"}, only_evaluate_locally=True
+        "distinct_id_random_22",
+        person_properties={"$geoip_city_name": "Sydney"},
+        only_evaluate_locally=True,
     )
 )
 print(posthog.get_remote_config_payload("encrypted_payload_flag_key"))
 
 
+# You can add tags to a context, and these are automatically added to any events (including exceptions) captured
+# within that context.
+
+# You can enter a new context using a with statement. Any exceptions thrown in the context will be captured,
+# and tagged with the context tags. Other events captured will also be tagged with the context tags. By default,
+# the new context inherits tags from the parent context.
+with posthog.new_context():
+    posthog.tag("transaction_id", "abc123")
+    posthog.tag("some_arbitrary_value", {"tags": "can be dicts"})
+
+    # This event will be captured with the tags set above
+    posthog.capture("order_processed")
+    # This exception will be captured with the tags set above
+    raise Exception("Order processing failed")
+
+
+# Use fresh=True to start with a clean context (no inherited tags)
+with posthog.new_context(fresh=True):
+    posthog.tag("session_id", "xyz789")
+    # Only session_id tag will be present, no inherited tags
+    raise Exception("Session handling failed")
+
+
+# You can also use the `@posthog.scoped()` decorator to enter a new context.
+# By default, it inherits tags from the parent context
+@posthog.scoped()
+def process_order(order_id):
+    posthog.tag("order_id", order_id)
+    # Exception will be captured and tagged automatically
+    raise Exception("Order processing failed")
+
+
+# Use fresh=True to start with a clean context (no inherited tags)
+@posthog.scoped(fresh=True)
+def process_payment(payment_id):
+    posthog.tag("payment_id", payment_id)
+    # Only payment_id tag will be present, no inherited tags
+    raise Exception("Payment processing failed")
+
+
 posthog.shutdown()

mypy-baseline.txt

@@ -1,7 +1,6 @@
 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]
@@ -38,4 +37,4 @@ 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]
+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

@@ -9,7 +9,7 @@ check_untyped_defs = True
 warn_unreachable = True
 strict_equality = True
 ignore_missing_imports = True
-exclude = env/.*|venv/.*
+exclude = env/.*|venv/.*|build/.*
 
 [mypy-django.*]
 ignore_missing_imports = True

posthog/__init__.py

@@ -4,11 +4,19 @@ 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.scopes import clear_tags, get_tags, new_context, scoped, tag
 from posthog.types import FeatureFlag, FlagsAndPayloads
 from posthog.version import VERSION
 
 __version__ = VERSION
 
+"""Context management."""
+new_context = new_context
+tag = tag
+get_tags = get_tags
+clear_tags = clear_tags
+tracked = scoped
+
 """Settings."""
 api_key = None  # type: Optional[str]
 host = None  # type: Optional[str]
@@ -315,7 +323,7 @@ def capture_exception(
     timestamp=None,  # type: Optional[datetime.datetime]
     uuid=None,  # type: Optional[str]
     groups=None,  # type: Optional[Dict]
-    **kwargs
+    **kwargs,
 ):
     # type: (...) -> Tuple[bool, dict]
     """
@@ -358,7 +366,7 @@ def capture_exception(
         timestamp=timestamp,
         uuid=uuid,
         groups=groups,
-        **kwargs
+        **kwargs,
     )
 
 
@@ -572,8 +580,7 @@ def shutdown():
     _proxy("join")
 
 
-def _proxy(method, *args, **kwargs):
-    """Create an analytics client if one doesn't exist and send to it."""
+def setup():
     global default_client
     if not default_client:
         default_client = Client(
@@ -602,6 +609,11 @@ def _proxy(method, *args, **kwargs):
     default_client.disabled = disabled
     default_client.debug = debug
 
+
+def _proxy(method, *args, **kwargs):
+    """Create an analytics client if one doesn't exist and send to it."""
+    setup()
+
     fn = getattr(default_client, method)
     return fn(*args, **kwargs)
 

posthog/ai/anthropic/__init__.py

@@ -1,6 +1,11 @@
 from .anthropic import Anthropic
 from .anthropic_async import AsyncAnthropic
-from .anthropic_providers import AnthropicBedrock, AnthropicVertex, AsyncAnthropicBedrock, AsyncAnthropicVertex
+from .anthropic_providers import (
+    AnthropicBedrock,
+    AnthropicVertex,
+    AsyncAnthropicBedrock,
+    AsyncAnthropicVertex,
+)
 
 __all__ = [
     "Anthropic",

posthog/ai/anthropic/anthropic.py

@@ -2,13 +2,20 @@ try:
     import anthropic
     from anthropic.resources import Messages
 except ImportError:
-    raise ModuleNotFoundError("Please install the Anthropic SDK to use this feature: 'pip install anthropic'")
+    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.ai.utils import (
+    call_llm_and_track_usage,
+    get_model_params,
+    merge_system_prompt,
+    with_privacy_mode,
+)
 from posthog.client import Client as PostHogClient
 
 
@@ -186,8 +193,12 @@ class WrappedMessages(Messages):
             "$ai_http_status": 200,
             "$ai_input_tokens": usage_stats.get("input_tokens", 0),
             "$ai_output_tokens": usage_stats.get("output_tokens", 0),
-            "$ai_cache_read_input_tokens": usage_stats.get("cache_read_input_tokens", 0),
-            "$ai_cache_creation_input_tokens": usage_stats.get("cache_creation_input_tokens", 0),
+            "$ai_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/ai/anthropic/anthropic_async.py

@@ -2,13 +2,20 @@ try:
     import anthropic
     from anthropic.resources import AsyncMessages
 except ImportError:
-    raise ModuleNotFoundError("Please install the Anthropic SDK to use this feature: 'pip install anthropic'")
+    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.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
 
 
@@ -186,8 +193,12 @@ class AsyncWrappedMessages(AsyncMessages):
             "$ai_http_status": 200,
             "$ai_input_tokens": usage_stats.get("input_tokens", 0),
             "$ai_output_tokens": usage_stats.get("output_tokens", 0),
-            "$ai_cache_read_input_tokens": usage_stats.get("cache_read_input_tokens", 0),
-            "$ai_cache_creation_input_tokens": usage_stats.get("cache_creation_input_tokens", 0),
+            "$ai_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/ai/anthropic/anthropic_providers.py

@@ -1,7 +1,9 @@
 try:
     import anthropic
 except ImportError:
-    raise ModuleNotFoundError("Please install the Anthropic SDK to use this feature: 'pip install anthropic'")
+    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

posthog/ai/gemini/__init__.py

@@ -0,0 +1,11 @@
+from .gemini import Client
+
+
+# Create a genai-like module for perfect drop-in replacement
+class _GenAI:
+    Client = Client
+
+
+genai = _GenAI()
+
+__all__ = ["Client", "genai"]

posthog/ai/gemini/gemini.py

@@ -0,0 +1,366 @@
+import os
+import time
+import uuid
+from typing import Any, Dict, Optional
+
+try:
+    from google import genai
+except ImportError:
+    raise ModuleNotFoundError(
+        "Please install the Google Gemini SDK to use this feature: 'pip install google-genai'"
+    )
+
+from posthog.ai.utils import (
+    call_llm_and_track_usage,
+    get_model_params,
+    with_privacy_mode,
+)
+from posthog.client import Client as PostHogClient
+
+
+class Client:
+    """
+    A drop-in replacement for genai.Client that automatically sends LLM usage events to PostHog.
+
+    Usage:
+        client = Client(
+            api_key="your_api_key",
+            posthog_client=posthog_client,
+            posthog_distinct_id="default_user",  # Optional defaults
+            posthog_properties={"team": "ai"}    # Optional defaults
+        )
+        response = client.models.generate_content(
+            model="gemini-2.0-flash",
+            contents=["Hello world"],
+            posthog_distinct_id="specific_user"  # Override default
+        )
+    """
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        posthog_client: Optional[PostHogClient] = None,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: bool = False,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            api_key: Google AI API key. If not provided, will use GOOGLE_API_KEY or API_KEY environment variable
+            posthog_client: PostHog client for tracking usage
+            posthog_distinct_id: Default distinct ID for all calls (can be overridden per call)
+            posthog_properties: Default properties for all calls (can be overridden per call)
+            posthog_privacy_mode: Default privacy mode for all calls (can be overridden per call)
+            posthog_groups: Default groups for all calls (can be overridden per call)
+            **kwargs: Additional arguments (for future compatibility)
+        """
+        if posthog_client is None:
+            raise ValueError("posthog_client is required for PostHog tracking")
+
+        self.models = Models(
+            api_key=api_key,
+            posthog_client=posthog_client,
+            posthog_distinct_id=posthog_distinct_id,
+            posthog_properties=posthog_properties,
+            posthog_privacy_mode=posthog_privacy_mode,
+            posthog_groups=posthog_groups,
+            **kwargs,
+        )
+
+
+class Models:
+    """
+    Models interface that mimics genai.Client().models with PostHog tracking.
+    """
+
+    _ph_client: PostHogClient  # Not None after __init__ validation
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        posthog_client: Optional[PostHogClient] = None,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: bool = False,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            api_key: Google AI API key. If not provided, will use GOOGLE_API_KEY or API_KEY environment variable
+            posthog_client: PostHog client for tracking usage
+            posthog_distinct_id: Default distinct ID for all calls
+            posthog_properties: Default properties for all calls
+            posthog_privacy_mode: Default privacy mode for all calls
+            posthog_groups: Default groups for all calls
+            **kwargs: Additional arguments (for future compatibility)
+        """
+        if posthog_client is None:
+            raise ValueError("posthog_client is required for PostHog tracking")
+
+        self._ph_client = posthog_client
+
+        # Store default PostHog settings
+        self._default_distinct_id = posthog_distinct_id
+        self._default_properties = posthog_properties or {}
+        self._default_privacy_mode = posthog_privacy_mode
+        self._default_groups = posthog_groups
+
+        # Handle API key - try parameter first, then environment variables
+        if api_key is None:
+            api_key = os.environ.get("GOOGLE_API_KEY") or os.environ.get("API_KEY")
+
+        if api_key is None:
+            raise ValueError(
+                "API key must be provided either as parameter or via GOOGLE_API_KEY/API_KEY environment variable"
+            )
+
+        self._client = genai.Client(api_key=api_key)
+        self._base_url = "https://generativelanguage.googleapis.com"
+
+    def _merge_posthog_params(
+        self,
+        call_distinct_id: Optional[str],
+        call_trace_id: Optional[str],
+        call_properties: Optional[Dict[str, Any]],
+        call_privacy_mode: Optional[bool],
+        call_groups: Optional[Dict[str, Any]],
+    ):
+        """Merge call-level PostHog parameters with client defaults."""
+        # Use call-level values if provided, otherwise fall back to defaults
+        distinct_id = (
+            call_distinct_id
+            if call_distinct_id is not None
+            else self._default_distinct_id
+        )
+        privacy_mode = (
+            call_privacy_mode
+            if call_privacy_mode is not None
+            else self._default_privacy_mode
+        )
+        groups = call_groups if call_groups is not None else self._default_groups
+
+        # Merge properties: default properties + call properties (call properties override)
+        properties = dict(self._default_properties)
+        if call_properties:
+            properties.update(call_properties)
+
+        if call_trace_id is None:
+            call_trace_id = str(uuid.uuid4())
+
+        return distinct_id, call_trace_id, properties, privacy_mode, groups
+
+    def generate_content(
+        self,
+        model: str,
+        contents,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: Optional[bool] = None,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        """
+        Generate content using Gemini's API while tracking usage in PostHog.
+
+        This method signature exactly matches genai.Client().models.generate_content()
+        with additional PostHog tracking parameters.
+
+        Args:
+            model: The model to use (e.g., 'gemini-2.0-flash')
+            contents: The input content for generation
+            posthog_distinct_id: ID to associate with the usage event (overrides client default)
+            posthog_trace_id: Trace UUID for linking events (auto-generated if not provided)
+            posthog_properties: Extra properties to include in the event (merged with client defaults)
+            posthog_privacy_mode: Whether to redact sensitive information (overrides client default)
+            posthog_groups: Group analytics properties (overrides client default)
+            **kwargs: Arguments passed to Gemini's generate_content
+        """
+        # Merge PostHog parameters
+        distinct_id, trace_id, properties, privacy_mode, groups = (
+            self._merge_posthog_params(
+                posthog_distinct_id,
+                posthog_trace_id,
+                posthog_properties,
+                posthog_privacy_mode,
+                posthog_groups,
+            )
+        )
+
+        kwargs_with_contents = {"model": model, "contents": contents, **kwargs}
+
+        return call_llm_and_track_usage(
+            distinct_id,
+            self._ph_client,
+            "gemini",
+            trace_id,
+            properties,
+            privacy_mode,
+            groups,
+            self._base_url,
+            self._client.models.generate_content,
+            **kwargs_with_contents,
+        )
+
+    def _generate_content_streaming(
+        self,
+        model: str,
+        contents,
+        distinct_id: Optional[str],
+        trace_id: Optional[str],
+        properties: Optional[Dict[str, Any]],
+        privacy_mode: bool,
+        groups: Optional[Dict[str, Any]],
+        **kwargs: Any,
+    ):
+        start_time = time.time()
+        usage_stats: Dict[str, int] = {"input_tokens": 0, "output_tokens": 0}
+        accumulated_content = []
+
+        kwargs_without_stream = {"model": model, "contents": contents, **kwargs}
+        response = self._client.models.generate_content_stream(**kwargs_without_stream)
+
+        def generator():
+            nonlocal usage_stats
+            nonlocal accumulated_content  # noqa: F824
+            try:
+                for chunk in response:
+                    if hasattr(chunk, "usage_metadata") and chunk.usage_metadata:
+                        usage_stats = {
+                            "input_tokens": getattr(
+                                chunk.usage_metadata, "prompt_token_count", 0
+                            ),
+                            "output_tokens": getattr(
+                                chunk.usage_metadata, "candidates_token_count", 0
+                            ),
+                        }
+
+                    if hasattr(chunk, "text") and chunk.text:
+                        accumulated_content.append(chunk.text)
+
+                    yield chunk
+
+            finally:
+                end_time = time.time()
+                latency = end_time - start_time
+                output = "".join(accumulated_content)
+
+                self._capture_streaming_event(
+                    model,
+                    contents,
+                    distinct_id,
+                    trace_id,
+                    properties,
+                    privacy_mode,
+                    groups,
+                    kwargs,
+                    usage_stats,
+                    latency,
+                    output,
+                )
+
+        return generator()
+
+    def _capture_streaming_event(
+        self,
+        model: str,
+        contents,
+        distinct_id: Optional[str],
+        trace_id: Optional[str],
+        properties: Optional[Dict[str, Any]],
+        privacy_mode: bool,
+        groups: Optional[Dict[str, Any]],
+        kwargs: Dict[str, Any],
+        usage_stats: Dict[str, int],
+        latency: float,
+        output: str,
+    ):
+        if trace_id is None:
+            trace_id = str(uuid.uuid4())
+
+        event_properties = {
+            "$ai_provider": "gemini",
+            "$ai_model": model,
+            "$ai_model_parameters": get_model_params(kwargs),
+            "$ai_input": with_privacy_mode(
+                self._ph_client,
+                privacy_mode,
+                self._format_input(contents),
+            ),
+            "$ai_output_choices": with_privacy_mode(
+                self._ph_client,
+                privacy_mode,
+                [{"content": output, "role": "assistant"}],
+            ),
+            "$ai_http_status": 200,
+            "$ai_input_tokens": usage_stats.get("input_tokens", 0),
+            "$ai_output_tokens": usage_stats.get("output_tokens", 0),
+            "$ai_latency": latency,
+            "$ai_trace_id": trace_id,
+            "$ai_base_url": self._base_url,
+            **(properties or {}),
+        }
+
+        if distinct_id is None:
+            event_properties["$process_person_profile"] = False
+
+        if hasattr(self._ph_client, "capture"):
+            self._ph_client.capture(
+                distinct_id=distinct_id,
+                event="$ai_generation",
+                properties=event_properties,
+                groups=groups,
+            )
+
+    def _format_input(self, contents):
+        """Format input contents for PostHog tracking"""
+        if isinstance(contents, str):
+            return [{"role": "user", "content": contents}]
+        elif isinstance(contents, list):
+            formatted = []
+            for item in contents:
+                if isinstance(item, str):
+                    formatted.append({"role": "user", "content": item})
+                elif hasattr(item, "text"):
+                    formatted.append({"role": "user", "content": item.text})
+                else:
+                    formatted.append({"role": "user", "content": str(item)})
+            return formatted
+        else:
+            return [{"role": "user", "content": str(contents)}]
+
+    def generate_content_stream(
+        self,
+        model: str,
+        contents,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: Optional[bool] = None,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        # Merge PostHog parameters
+        distinct_id, trace_id, properties, privacy_mode, groups = (
+            self._merge_posthog_params(
+                posthog_distinct_id,
+                posthog_trace_id,
+                posthog_properties,
+                posthog_privacy_mode,
+                posthog_groups,
+            )
+        )
+
+        return self._generate_content_streaming(
+            model,
+            contents,
+            distinct_id,
+            trace_id,
+            properties,
+            privacy_mode,
+            groups,
+            **kwargs,
+        )

posthog/ai/langchain/callbacks.py

@@ -1,7 +1,9 @@
 try:
     import langchain  # noqa: F401
 except ImportError:
-    raise ModuleNotFoundError("Please install LangChain to use this feature: 'pip install langchain'")
+    raise ModuleNotFoundError(
+        "Please install LangChain to use this feature: 'pip install langchain'"
+    )
 
 import logging
 import time
@@ -21,7 +23,14 @@ 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.messages import (
+    AIMessage,
+    BaseMessage,
+    FunctionMessage,
+    HumanMessage,
+    SystemMessage,
+    ToolMessage,
+)
 from langchain_core.outputs import ChatGeneration, LLMResult
 from pydantic import BaseModel
 
@@ -143,7 +152,9 @@ class CallbackHandler(BaseCallbackHandler):
     ):
         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)
+        self._set_trace_or_span_metadata(
+            serialized, inputs, run_id, parent_run_id, **kwargs
+        )
 
     def on_chain_end(
         self,
@@ -176,9 +187,13 @@ class CallbackHandler(BaseCallbackHandler):
         parent_run_id: Optional[UUID] = None,
         **kwargs,
     ):
-        self._log_debug_event("on_chat_model_start", run_id, parent_run_id, messages=messages)
+        self._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]
+        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(
@@ -216,7 +231,9 @@ class CallbackHandler(BaseCallbackHandler):
         """
         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._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(
@@ -240,9 +257,13 @@ class CallbackHandler(BaseCallbackHandler):
         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._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)
+        self._set_trace_or_span_metadata(
+            serialized, input_str, run_id, parent_run_id, **kwargs
+        )
 
     def on_tool_end(
         self,
@@ -279,7 +300,9 @@ class CallbackHandler(BaseCallbackHandler):
     ) -> 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)
+        self._set_trace_or_span_metadata(
+            serialized, query, run_id, parent_run_id, **kwargs
+        )
 
     def on_retriever_end(
         self,
@@ -289,7 +312,9 @@ class CallbackHandler(BaseCallbackHandler):
         parent_run_id: Optional[UUID] = None,
         **kwargs: Any,
     ):
-        self._log_debug_event("on_retriever_end", run_id, parent_run_id, documents=documents)
+        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(
@@ -364,7 +389,9 @@ class CallbackHandler(BaseCallbackHandler):
     ):
         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)
+        self._runs[run_id] = SpanMetadata(
+            name=run_name, input=input, start_time=time.time(), end_time=None
+        )
 
     def _set_llm_metadata(
         self,
@@ -376,7 +403,9 @@ class CallbackHandler(BaseCallbackHandler):
         **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)
+        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"):
@@ -410,7 +439,9 @@ class CallbackHandler(BaseCallbackHandler):
             return run_id
         return trace_id
 
-    def _get_parent_run_id(self, trace_id: Any, run_id: UUID, parent_run_id: Optional[UUID]):
+    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.
         """
@@ -418,14 +449,18 @@ class CallbackHandler(BaseCallbackHandler):
             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):
+    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.")
+            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,
@@ -446,7 +481,9 @@ class CallbackHandler(BaseCallbackHandler):
         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_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,
@@ -460,7 +497,9 @@ class CallbackHandler(BaseCallbackHandler):
             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)
+            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
@@ -484,7 +523,9 @@ class CallbackHandler(BaseCallbackHandler):
         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.")
+            log.warning(
+                f"Run {run_id} is not a generation, but attempted to be captured as a generation."
+            )
             return
         self._capture_generation(
             trace_id,
@@ -540,8 +581,12 @@ class CallbackHandler(BaseCallbackHandler):
                     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)
+                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)
@@ -631,7 +676,9 @@ def _parse_usage_model(
         if model_key in usage:
             captured_count = usage[model_key]
             final_count = (
-                sum(captured_count) if isinstance(captured_count, list) else captured_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
@@ -656,8 +703,12 @@ def _parse_usage(response: LLMResult):
                 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"])
+                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", {})
@@ -669,13 +720,19 @@ def _parse_usage(response: LLMResult):
                     else None
                 )
                 bedrock_titan_usage = (
-                    response_metadata.get("amazon-bedrock-invocationMetrics", None)  # for Bedrock-Titan
+                    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
+                ollama_usage = getattr(
+                    message_chunk, "usage_metadata", None
+                )  # for Ollama
 
-                chunk_usage = bedrock_anthropic_usage or bedrock_titan_usage or ollama_usage
+                chunk_usage = (
+                    bedrock_anthropic_usage or bedrock_titan_usage or ollama_usage
+                )
                 if chunk_usage:
                     llm_usage = _parse_usage_model(chunk_usage)
                     break
@@ -691,7 +748,9 @@ def _get_http_status(error: BaseException) -> int:
     return status_code
 
 
-def _get_langchain_run_name(serialized: Optional[Dict[str, Any]], **kwargs: Any) -> Optional[str]:
+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:

posthog/ai/openai/openai.py

@@ -4,11 +4,16 @@ 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'")
+    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.ai.utils import (
+    call_llm_and_track_usage,
+    get_model_params,
+    with_privacy_mode,
+)
 from posthog.client import Client as PostHogClient
 
 
@@ -29,14 +34,37 @@ class OpenAI(openai.OpenAI):
         """
         super().__init__(**kwargs)
         self._ph_client = posthog_client
-        self.chat = WrappedChat(self)
-        self.embeddings = WrappedEmbeddings(self)
-        self.beta = WrappedBeta(self)
-        self.responses = WrappedResponses(self)
+
+        # Store original objects after parent initialization (only if they exist)
+        self._original_chat = getattr(self, "chat", None)
+        self._original_embeddings = getattr(self, "embeddings", None)
+        self._original_beta = getattr(self, "beta", None)
+        self._original_responses = getattr(self, "responses", None)
+
+        # Replace with wrapped versions (only if originals exist)
+        if self._original_chat is not None:
+            self.chat = WrappedChat(self, self._original_chat)
+
+        if self._original_embeddings is not None:
+            self.embeddings = WrappedEmbeddings(self, self._original_embeddings)
+
+        if self._original_beta is not None:
+            self.beta = WrappedBeta(self, self._original_beta)
+
+        if self._original_responses is not None:
+            self.responses = WrappedResponses(self, self._original_responses)
 
 
-class WrappedResponses(openai.resources.responses.Responses):
-    _client: OpenAI
+class WrappedResponses:
+    """Wrapper for OpenAI responses that tracks usage in PostHog."""
+
+    def __init__(self, client: OpenAI, original_responses):
+        self._client = client
+        self._original = original_responses
+
+    def __getattr__(self, name):
+        """Fallback to original responses object for any methods we don't explicitly handle."""
+        return getattr(self._original, name)
 
     def create(
         self,
@@ -69,7 +97,7 @@ class WrappedResponses(openai.resources.responses.Responses):
             posthog_privacy_mode,
             posthog_groups,
             self._client.base_url,
-            super().create,
+            self._original.create,
             **kwargs,
         )
 
@@ -85,7 +113,7 @@ class WrappedResponses(openai.resources.responses.Responses):
         start_time = time.time()
         usage_stats: Dict[str, int] = {}
         final_content = []
-        response = super().create(**kwargs)
+        response = self._original.create(**kwargs)
 
         def generator():
             nonlocal usage_stats
@@ -112,12 +140,16 @@ class WrappedResponses(openai.resources.responses.Responses):
                         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
+                            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
+                            usage_stats["cache_read_input_tokens"] = (
+                                chunk.usage.input_tokens_details.cached_tokens
+                            )
 
                     yield chunk
 
@@ -159,7 +191,9 @@ class WrappedResponses(openai.resources.responses.Responses):
             "$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_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,
@@ -168,7 +202,9 @@ class WrappedResponses(openai.resources.responses.Responses):
             "$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_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,
@@ -195,16 +231,32 @@ class WrappedResponses(openai.resources.responses.Responses):
             )
 
 
-class WrappedChat(openai.resources.chat.Chat):
-    _client: OpenAI
+class WrappedChat:
+    """Wrapper for OpenAI chat that tracks usage in PostHog."""
+
+    def __init__(self, client: OpenAI, original_chat):
+        self._client = client
+        self._original = original_chat
+
+    def __getattr__(self, name):
+        """Fallback to original chat object for any methods we don't explicitly handle."""
+        return getattr(self._original, name)
 
     @property
     def completions(self):
-        return WrappedCompletions(self._client)
+        return WrappedCompletions(self._client, self._original.completions)
 
 
-class WrappedCompletions(openai.resources.chat.completions.Completions):
-    _client: OpenAI
+class WrappedCompletions:
+    """Wrapper for OpenAI chat completions that tracks usage in PostHog."""
+
+    def __init__(self, client: OpenAI, original_completions):
+        self._client = client
+        self._original = original_completions
+
+    def __getattr__(self, name):
+        """Fallback to original completions object for any methods we don't explicitly handle."""
+        return getattr(self._original, name)
 
     def create(
         self,
@@ -237,7 +289,7 @@ class WrappedCompletions(openai.resources.chat.completions.Completions):
             posthog_privacy_mode,
             posthog_groups,
             self._client.base_url,
-            super().create,
+            self._original.create,
             **kwargs,
         )
 
@@ -257,7 +309,7 @@ class WrappedCompletions(openai.resources.chat.completions.Completions):
         if "stream_options" not in kwargs:
             kwargs["stream_options"] = {}
         kwargs["stream_options"]["include_usage"] = True
-        response = super().create(**kwargs)
+        response = self._original.create(**kwargs)
 
         def generator():
             nonlocal usage_stats
@@ -280,14 +332,22 @@ class WrappedCompletions(openai.resources.chat.completions.Completions):
                         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
+                            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
+                            usage_stats["reasoning_tokens"] = (
+                                chunk.usage.output_tokens_details.reasoning_tokens
+                            )
 
-                    if hasattr(chunk, "choices") and chunk.choices and len(chunk.choices) > 0:
+                    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:
@@ -302,8 +362,14 @@ class WrappedCompletions(openai.resources.chat.completions.Completions):
                                     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
+                                    if hasattr(tool_call, "function") and hasattr(
+                                        tool_call.function, "arguments"
+                                    ):
+                                        accumulated_tools[
+                                            index
+                                        ].function.arguments += (
+                                            tool_call.function.arguments
+                                        )
 
                     yield chunk
 
@@ -347,7 +413,9 @@ class WrappedCompletions(openai.resources.chat.completions.Completions):
             "$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_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,
@@ -356,7 +424,9 @@ class WrappedCompletions(openai.resources.chat.completions.Completions):
             "$ai_http_status": 200,
             "$ai_input_tokens": usage_stats.get("prompt_tokens", 0),
             "$ai_output_tokens": usage_stats.get("completion_tokens", 0),
-            "$ai_cache_read_input_tokens": usage_stats.get("cache_read_input_tokens", 0),
+            "$ai_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,
@@ -383,8 +453,16 @@ class WrappedCompletions(openai.resources.chat.completions.Completions):
             )
 
 
-class WrappedEmbeddings(openai.resources.embeddings.Embeddings):
-    _client: OpenAI
+class WrappedEmbeddings:
+    """Wrapper for OpenAI embeddings that tracks usage in PostHog."""
+
+    def __init__(self, client: OpenAI, original_embeddings):
+        self._client = client
+        self._original = original_embeddings
+
+    def __getattr__(self, name):
+        """Fallback to original embeddings object for any methods we don't explicitly handle."""
+        return getattr(self._original, name)
 
     def create(
         self,
@@ -402,6 +480,8 @@ class WrappedEmbeddings(openai.resources.embeddings.Embeddings):
             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 anonymize the input and output.
+            posthog_groups: Optional dictionary of groups to associate with the event.
             **kwargs: Any additional parameters for the OpenAI Embeddings API.
 
         Returns:
@@ -411,7 +491,7 @@ class WrappedEmbeddings(openai.resources.embeddings.Embeddings):
             posthog_trace_id = str(uuid.uuid4())
 
         start_time = time.time()
-        response = super().create(**kwargs)
+        response = self._original.create(**kwargs)
         end_time = time.time()
 
         # Extract usage statistics if available
@@ -428,7 +508,9 @@ class WrappedEmbeddings(openai.resources.embeddings.Embeddings):
         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_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,
@@ -452,24 +534,48 @@ class WrappedEmbeddings(openai.resources.embeddings.Embeddings):
         return response
 
 
-class WrappedBeta(openai.resources.beta.Beta):
-    _client: OpenAI
+class WrappedBeta:
+    """Wrapper for OpenAI beta features that tracks usage in PostHog."""
+
+    def __init__(self, client: OpenAI, original_beta):
+        self._client = client
+        self._original = original_beta
+
+    def __getattr__(self, name):
+        """Fallback to original beta object for any methods we don't explicitly handle."""
+        return getattr(self._original, name)
 
     @property
     def chat(self):
-        return WrappedBetaChat(self._client)
+        return WrappedBetaChat(self._client, self._original.chat)
 
 
-class WrappedBetaChat(openai.resources.beta.chat.Chat):
-    _client: OpenAI
+class WrappedBetaChat:
+    """Wrapper for OpenAI beta chat that tracks usage in PostHog."""
+
+    def __init__(self, client: OpenAI, original_beta_chat):
+        self._client = client
+        self._original = original_beta_chat
+
+    def __getattr__(self, name):
+        """Fallback to original beta chat object for any methods we don't explicitly handle."""
+        return getattr(self._original, name)
 
     @property
     def completions(self):
-        return WrappedBetaCompletions(self._client)
+        return WrappedBetaCompletions(self._client, self._original.completions)
 
 
-class WrappedBetaCompletions(openai.resources.beta.chat.completions.Completions):
-    _client: OpenAI
+class WrappedBetaCompletions:
+    """Wrapper for OpenAI beta chat completions that tracks usage in PostHog."""
+
+    def __init__(self, client: OpenAI, original_beta_completions):
+        self._client = client
+        self._original = original_beta_completions
+
+    def __getattr__(self, name):
+        """Fallback to original beta completions object for any methods we don't explicitly handle."""
+        return getattr(self._original, name)
 
     def parse(
         self,
@@ -489,6 +595,6 @@ class WrappedBetaCompletions(openai.resources.beta.chat.completions.Completions)
             posthog_privacy_mode,
             posthog_groups,
             self._client.base_url,
-            super().parse,
+            self._original.parse,
             **kwargs,
         )

posthog/ai/openai/openai_async.py

@@ -4,11 +4,16 @@ 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'")
+    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.ai.utils import (
+    call_llm_and_track_usage_async,
+    get_model_params,
+    with_privacy_mode,
+)
 from posthog.client import Client as PostHogClient
 
 
@@ -23,19 +28,43 @@ class AsyncOpenAI(openai.AsyncOpenAI):
         """
         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").
+            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)
+
+        # Store original objects after parent initialization (only if they exist)
+        self._original_chat = getattr(self, "chat", None)
+        self._original_embeddings = getattr(self, "embeddings", None)
+        self._original_beta = getattr(self, "beta", None)
+        self._original_responses = getattr(self, "responses", None)
+
+        # Replace with wrapped versions (only if originals exist)
+        if self._original_chat is not None:
+            self.chat = WrappedChat(self, self._original_chat)
+
+        if self._original_embeddings is not None:
+            self.embeddings = WrappedEmbeddings(self, self._original_embeddings)
+
+        if self._original_beta is not None:
+            self.beta = WrappedBeta(self, self._original_beta)
+
+        if self._original_responses is not None:
+            self.responses = WrappedResponses(self, self._original_responses)
 
 
-class WrappedResponses(openai.resources.responses.Responses):
-    _client: AsyncOpenAI
+class WrappedResponses:
+    """Async wrapper for OpenAI responses that tracks usage in PostHog."""
+
+    def __init__(self, client: AsyncOpenAI, original_responses):
+        self._client = client
+        self._original = original_responses
+
+    def __getattr__(self, name):
+        """Fallback to original responses object for any methods we don't explicitly handle."""
+        return getattr(self._original, name)
 
     async def create(
         self,
@@ -68,7 +97,7 @@ class WrappedResponses(openai.resources.responses.Responses):
             posthog_privacy_mode,
             posthog_groups,
             self._client.base_url,
-            super().create,
+            self._original.create,
             **kwargs,
         )
 
@@ -84,7 +113,7 @@ class WrappedResponses(openai.resources.responses.Responses):
         start_time = time.time()
         usage_stats: Dict[str, int] = {}
         final_content = []
-        response = await super().create(**kwargs)
+        response = await self._original.create(**kwargs)
 
         async def async_generator():
             nonlocal usage_stats
@@ -111,12 +140,16 @@ class WrappedResponses(openai.resources.responses.Responses):
                         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
+                            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
+                            usage_stats["cache_read_input_tokens"] = (
+                                chunk.usage.input_tokens_details.cached_tokens
+                            )
 
                     yield chunk
 
@@ -158,7 +191,9 @@ class WrappedResponses(openai.resources.responses.Responses):
             "$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_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,
@@ -167,7 +202,9 @@ class WrappedResponses(openai.resources.responses.Responses):
             "$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_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,
@@ -186,7 +223,7 @@ class WrappedResponses(openai.resources.responses.Responses):
             event_properties["$process_person_profile"] = False
 
         if hasattr(self._client._ph_client, "capture"):
-            await 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,
@@ -194,16 +231,32 @@ class WrappedResponses(openai.resources.responses.Responses):
             )
 
 
-class WrappedChat(openai.resources.chat.AsyncChat):
-    _client: AsyncOpenAI
+class WrappedChat:
+    """Async wrapper for OpenAI chat that tracks usage in PostHog."""
+
+    def __init__(self, client: AsyncOpenAI, original_chat):
+        self._client = client
+        self._original = original_chat
+
+    def __getattr__(self, name):
+        """Fallback to original chat object for any methods we don't explicitly handle."""
+        return getattr(self._original, name)
 
     @property
     def completions(self):
-        return WrappedCompletions(self._client)
+        return WrappedCompletions(self._client, self._original.completions)
 
 
-class WrappedCompletions(openai.resources.chat.completions.AsyncCompletions):
-    _client: AsyncOpenAI
+class WrappedCompletions:
+    """Async wrapper for OpenAI chat completions that tracks usage in PostHog."""
+
+    def __init__(self, client: AsyncOpenAI, original_completions):
+        self._client = client
+        self._original = original_completions
+
+    def __getattr__(self, name):
+        """Fallback to original completions object for any methods we don't explicitly handle."""
+        return getattr(self._original, name)
 
     async def create(
         self,
@@ -237,7 +290,7 @@ class WrappedCompletions(openai.resources.chat.completions.AsyncCompletions):
             posthog_privacy_mode,
             posthog_groups,
             self._client.base_url,
-            super().create,
+            self._original.create,
             **kwargs,
         )
         return response
@@ -247,21 +300,25 @@ class WrappedCompletions(openai.resources.chat.completions.AsyncCompletions):
         posthog_distinct_id: Optional[str],
         posthog_trace_id: Optional[str],
         posthog_properties: Optional[Dict[str, Any]],
-        posthog_privacy_mode: bool = False,
-        posthog_groups: Optional[Dict[str, Any]] = None,
+        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 = await super().create(**kwargs)
+        response = await self._original.create(**kwargs)
 
         async def async_generator():
-            nonlocal usage_stats, accumulated_content, accumulated_tools  # noqa: F824
+            nonlocal usage_stats
+            nonlocal accumulated_content  # noqa: F824
+            nonlocal accumulated_tools  # noqa: F824
+
             try:
                 async for chunk in response:
                     if hasattr(chunk, "usage") and chunk.usage:
@@ -278,9 +335,22 @@ class WrappedCompletions(openai.resources.chat.completions.AsyncCompletions):
                         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
+                            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 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:
@@ -295,8 +365,14 @@ class WrappedCompletions(openai.resources.chat.completions.AsyncCompletions):
                                     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
+                                    if hasattr(tool_call, "function") and hasattr(
+                                        tool_call.function, "arguments"
+                                    ):
+                                        accumulated_tools[
+                                            index
+                                        ].function.arguments += (
+                                            tool_call.function.arguments
+                                        )
 
                     yield chunk
 
@@ -340,7 +416,9 @@ class WrappedCompletions(openai.resources.chat.completions.AsyncCompletions):
             "$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_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,
@@ -349,7 +427,10 @@ class WrappedCompletions(openai.resources.chat.completions.AsyncCompletions):
             "$ai_http_status": 200,
             "$ai_input_tokens": usage_stats.get("prompt_tokens", 0),
             "$ai_output_tokens": usage_stats.get("completion_tokens", 0),
-            "$ai_cache_read_input_tokens": usage_stats.get("cache_read_input_tokens", 0),
+            "$ai_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),
@@ -367,7 +448,7 @@ class WrappedCompletions(openai.resources.chat.completions.AsyncCompletions):
             event_properties["$process_person_profile"] = False
 
         if hasattr(self._client._ph_client, "capture"):
-            await 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,
@@ -375,8 +456,16 @@ class WrappedCompletions(openai.resources.chat.completions.AsyncCompletions):
             )
 
 
-class WrappedEmbeddings(openai.resources.embeddings.AsyncEmbeddings):
-    _client: AsyncOpenAI
+class WrappedEmbeddings:
+    """Async wrapper for OpenAI embeddings that tracks usage in PostHog."""
+
+    def __init__(self, client: AsyncOpenAI, original_embeddings):
+        self._client = client
+        self._original = original_embeddings
+
+    def __getattr__(self, name):
+        """Fallback to original embeddings object for any methods we don't explicitly handle."""
+        return getattr(self._original, name)
 
     async def create(
         self,
@@ -394,8 +483,8 @@ class WrappedEmbeddings(openai.resources.embeddings.AsyncEmbeddings):
             posthog_distinct_id: Optional ID to associate with the usage event.
             posthog_trace_id: Optional trace UUID for linking events.
             posthog_properties: Optional dictionary of extra properties to include in the event.
-            posthog_privacy_mode: Whether to store input and output in PostHog.
-            posthog_groups: Optional dictionary of groups to include in the event.
+            posthog_privacy_mode: Whether to anonymize the input and output.
+            posthog_groups: Optional dictionary of groups to associate with the event.
             **kwargs: Any additional parameters for the OpenAI Embeddings API.
 
         Returns:
@@ -405,7 +494,7 @@ class WrappedEmbeddings(openai.resources.embeddings.AsyncEmbeddings):
             posthog_trace_id = str(uuid.uuid4())
 
         start_time = time.time()
-        response = await super().create(**kwargs)
+        response = await self._original.create(**kwargs)
         end_time = time.time()
 
         # Extract usage statistics if available
@@ -422,7 +511,9 @@ class WrappedEmbeddings(openai.resources.embeddings.AsyncEmbeddings):
         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_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,
@@ -446,24 +537,48 @@ class WrappedEmbeddings(openai.resources.embeddings.AsyncEmbeddings):
         return response
 
 
-class WrappedBeta(openai.resources.beta.AsyncBeta):
-    _client: AsyncOpenAI
+class WrappedBeta:
+    """Async wrapper for OpenAI beta features that tracks usage in PostHog."""
+
+    def __init__(self, client: AsyncOpenAI, original_beta):
+        self._client = client
+        self._original = original_beta
+
+    def __getattr__(self, name):
+        """Fallback to original beta object for any methods we don't explicitly handle."""
+        return getattr(self._original, name)
 
     @property
     def chat(self):
-        return WrappedBetaChat(self._client)
+        return WrappedBetaChat(self._client, self._original.chat)
 
 
-class WrappedBetaChat(openai.resources.beta.chat.AsyncChat):
-    _client: AsyncOpenAI
+class WrappedBetaChat:
+    """Async wrapper for OpenAI beta chat that tracks usage in PostHog."""
+
+    def __init__(self, client: AsyncOpenAI, original_beta_chat):
+        self._client = client
+        self._original = original_beta_chat
+
+    def __getattr__(self, name):
+        """Fallback to original beta chat object for any methods we don't explicitly handle."""
+        return getattr(self._original, name)
 
     @property
     def completions(self):
-        return WrappedBetaCompletions(self._client)
+        return WrappedBetaCompletions(self._client, self._original.completions)
 
 
-class WrappedBetaCompletions(openai.resources.beta.chat.completions.AsyncCompletions):
-    _client: AsyncOpenAI
+class WrappedBetaCompletions:
+    """Async wrapper for OpenAI beta chat completions that tracks usage in PostHog."""
+
+    def __init__(self, client: AsyncOpenAI, original_beta_completions):
+        self._client = client
+        self._original = original_beta_completions
+
+    def __getattr__(self, name):
+        """Fallback to original beta completions object for any methods we don't explicitly handle."""
+        return getattr(self._original, name)
 
     async def parse(
         self,
@@ -483,6 +598,6 @@ class WrappedBetaCompletions(openai.resources.beta.chat.completions.AsyncComplet
             posthog_privacy_mode,
             posthog_groups,
             self._client.base_url,
-            super().parse,
+            self._original.parse,
             **kwargs,
         )

posthog/ai/openai/openai_providers.py

@@ -1,13 +1,20 @@
 try:
     import openai
-    import openai.resources
 except ImportError:
-    raise ModuleNotFoundError("Please install the Open AI SDK to use this feature: 'pip install openai'")
+    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 import (
+    WrappedBeta,
+    WrappedChat,
+    WrappedEmbeddings,
+    WrappedResponses,
+)
 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.ai.openai.openai_async import WrappedResponses as AsyncWrappedResponses
 from posthog.client import Client as PostHogClient
 
 
@@ -19,23 +26,70 @@ class AzureOpenAI(openai.AzureOpenAI):
     _ph_client: PostHogClient
 
     def __init__(self, posthog_client: PostHogClient, **kwargs):
+        """
+        Args:
+            api_key: Azure 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 Azure OpenAI (e.g. azure_endpoint="xxx").
+        """
         super().__init__(**kwargs)
         self._ph_client = posthog_client
-        self.chat = WrappedChat(self)
-        self.embeddings = WrappedEmbeddings(self)
-        self.beta = WrappedBeta(self)
+
+        # Store original objects after parent initialization (only if they exist)
+        self._original_chat = getattr(self, "chat", None)
+        self._original_embeddings = getattr(self, "embeddings", None)
+        self._original_beta = getattr(self, "beta", None)
+        self._original_responses = getattr(self, "responses", None)
+
+        # Replace with wrapped versions (only if originals exist)
+        if self._original_chat is not None:
+            self.chat = WrappedChat(self, self._original_chat)
+
+        if self._original_embeddings is not None:
+            self.embeddings = WrappedEmbeddings(self, self._original_embeddings)
+
+        if self._original_beta is not None:
+            self.beta = WrappedBeta(self, self._original_beta)
+
+        if self._original_responses is not None:
+            self.responses = WrappedResponses(self, self._original_responses)
 
 
 class AsyncAzureOpenAI(openai.AsyncAzureOpenAI):
     """
-    A wrapper around the Azure OpenAI SDK that automatically sends LLM usage events to PostHog.
+    An async wrapper around the Azure OpenAI SDK that automatically sends LLM usage events to PostHog.
     """
 
     _ph_client: PostHogClient
 
     def __init__(self, posthog_client: PostHogClient, **kwargs):
+        """
+        Args:
+            api_key: Azure 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 Azure OpenAI (e.g. azure_endpoint="xxx").
+        """
         super().__init__(**kwargs)
         self._ph_client = posthog_client
-        self.chat = AsyncWrappedChat(self)
-        self.embeddings = AsyncWrappedEmbeddings(self)
-        self.beta = AsyncWrappedBeta(self)
+
+        # Store original objects after parent initialization (only if they exist)
+        self._original_chat = getattr(self, "chat", None)
+        self._original_embeddings = getattr(self, "embeddings", None)
+        self._original_beta = getattr(self, "beta", None)
+        self._original_responses = getattr(self, "responses", None)
+
+        # Replace with wrapped versions (only if originals exist)
+        if self._original_chat is not None:
+            self.chat = AsyncWrappedChat(self, self._original_chat)
+
+        if self._original_embeddings is not None:
+            self.embeddings = AsyncWrappedEmbeddings(self, self._original_embeddings)
+
+        if self._original_beta is not None:
+            self.beta = AsyncWrappedBeta(self, self._original_beta)
+
+        # Only add responses if available (newer OpenAI versions)
+        if self._original_responses is not None:
+            self.responses = AsyncWrappedResponses(self, self._original_responses)

posthog/ai/utils.py

@@ -73,6 +73,23 @@ def get_usage(response, provider: str) -> Dict[str, Any]:
             "cache_read_input_tokens": cached_tokens,
             "reasoning_tokens": reasoning_tokens,
         }
+    elif provider == "gemini":
+        input_tokens = 0
+        output_tokens = 0
+
+        if hasattr(response, "usage_metadata") and response.usage_metadata:
+            input_tokens = getattr(response.usage_metadata, "prompt_token_count", 0)
+            output_tokens = getattr(
+                response.usage_metadata, "candidates_token_count", 0
+            )
+
+        return {
+            "input_tokens": input_tokens,
+            "output_tokens": output_tokens,
+            "cache_read_input_tokens": 0,
+            "cache_creation_input_tokens": 0,
+            "reasoning_tokens": 0,
+        }
     return {
         "input_tokens": 0,
         "output_tokens": 0,
@@ -93,6 +110,8 @@ def format_response(response, provider: str):
         return format_response_anthropic(response)
     elif provider == "openai":
         return format_response_openai(response)
+    elif provider == "gemini":
+        return format_response_gemini(response)
     return output
 
 
@@ -170,6 +189,40 @@ def format_response_openai(response):
     return output
 
 
+def format_response_gemini(response):
+    output = []
+    if hasattr(response, "candidates") and response.candidates:
+        for candidate in response.candidates:
+            if hasattr(candidate, "content") and candidate.content:
+                content_text = ""
+                if hasattr(candidate.content, "parts") and candidate.content.parts:
+                    for part in candidate.content.parts:
+                        if hasattr(part, "text") and part.text:
+                            content_text += part.text
+                if content_text:
+                    output.append(
+                        {
+                            "role": "assistant",
+                            "content": content_text,
+                        }
+                    )
+            elif hasattr(candidate, "text") and candidate.text:
+                output.append(
+                    {
+                        "role": "assistant",
+                        "content": candidate.text,
+                    }
+                )
+    elif hasattr(response, "text") and response.text:
+        output.append(
+            {
+                "role": "assistant",
+                "content": response.text,
+            }
+        )
+    return output
+
+
 def format_tool_calls(response, provider: str):
     if provider == "anthropic":
         if hasattr(response, "tools") and response.tools and len(response.tools) > 0:
@@ -186,7 +239,10 @@ def format_tool_calls(response, provider: str):
                 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:
+            if (
+                hasattr(response.choices[0], "tool_calls")
+                and response.choices[0].tool_calls
+            ):
                 return response.choices[0].tool_calls
     return None
 
@@ -198,6 +254,22 @@ def merge_system_prompt(kwargs: Dict[str, Any], provider: str):
         if kwargs.get("system") is None:
             return messages
         return [{"role": "system", "content": kwargs.get("system")}] + messages
+    elif provider == "gemini":
+        contents = kwargs.get("contents", [])
+        if isinstance(contents, str):
+            return [{"role": "user", "content": contents}]
+        elif isinstance(contents, list):
+            formatted = []
+            for item in contents:
+                if isinstance(item, str):
+                    formatted.append({"role": "user", "content": item})
+                elif hasattr(item, "text"):
+                    formatted.append({"role": "user", "content": item.text})
+                else:
+                    formatted.append({"role": "user", "content": str(item)})
+            return formatted
+        else:
+            return [{"role": "user", "content": str(contents)}]
 
     # For OpenAI, handle both Chat Completions and Responses API
     if kwargs.get("messages") is not None:
@@ -219,15 +291,21 @@ def merge_system_prompt(kwargs: Dict[str, Any], provider: str):
     # 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)
+        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')}"
+            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
+            messages = [
+                {"role": "system", "content": kwargs.get("instructions")}
+            ] + messages
 
     return messages
 
@@ -259,7 +337,9 @@ def call_llm_and_track_usage(
         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
+        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__(),
@@ -271,7 +351,10 @@ def call_llm_and_track_usage(
         if posthog_trace_id is None:
             posthog_trace_id = str(uuid.uuid4())
 
-        if response and hasattr(response, "usage"):
+        if response and (
+            hasattr(response, "usage")
+            or (provider == "gemini" and hasattr(response, "usage_metadata"))
+        ):
             usage = get_usage(response, provider)
 
         messages = merge_system_prompt(kwargs, provider)
@@ -296,15 +379,30 @@ def call_llm_and_track_usage(
 
         tool_calls = format_tool_calls(response, provider)
         if tool_calls:
-            event_properties["$ai_tools"] = with_privacy_mode(ph_client, posthog_privacy_mode, 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_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("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:
+        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:
@@ -354,7 +452,9 @@ async def call_llm_and_track_usage_async(
         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
+        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__(),
@@ -366,7 +466,10 @@ async def call_llm_and_track_usage_async(
         if posthog_trace_id is None:
             posthog_trace_id = str(uuid.uuid4())
 
-        if response and hasattr(response, "usage"):
+        if response and (
+            hasattr(response, "usage")
+            or (provider == "gemini" and hasattr(response, "usage_metadata"))
+        ):
             usage = get_usage(response, provider)
 
         messages = merge_system_prompt(kwargs, provider)
@@ -391,13 +494,25 @@ async def call_llm_and_track_usage_async(
 
         tool_calls = format_tool_calls(response, provider)
         if tool_calls:
-            event_properties["$ai_tools"] = with_privacy_mode(ph_client, posthog_privacy_mode, 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_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("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

posthog/client.py

@@ -16,7 +16,11 @@ from six import string_types
 
 from posthog.consumer import Consumer
 from posthog.exception_capture import ExceptionCapture
-from posthog.exception_utils import exc_info_from_error, exceptions_from_error_tuple, handle_in_app
+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 (
@@ -29,6 +33,7 @@ from posthog.request import (
     get,
     remote_config,
 )
+from posthog.scopes import get_tags
 from posthog.types import (
     FeatureFlag,
     FeatureFlagResult,
@@ -55,7 +60,9 @@ MAX_DICT_SIZE = 50_000
 
 # TODO: Get rid of these when you're done rolling out `/flags` to all customers
 ROLLOUT_PERCENTAGE = 1
-INCLUDED_HASHES = set({"bc94e67150c97dbcbf52549d50a7b80814841dbf"})  # this is PostHog's API key
+INCLUDED_HASHES = set(
+    {"bc94e67150c97dbcbf52549d50a7b80814841dbf"}
+)  # this is PostHog's API key
 # Explicitly excluding all the API tokens associated with the top 10 customers; we'll get to them soon, but don't want to rollout to them just yet
 EXCLUDED_HASHES = set(
     {
@@ -270,7 +277,9 @@ class Client(object):
         self.group_type_mapping = None
         self.cohorts = None
         self.poll_interval = poll_interval
-        self.feature_flags_request_timeout_seconds = feature_flags_request_timeout_seconds
+        self.feature_flags_request_timeout_seconds = (
+            feature_flags_request_timeout_seconds
+        )
         self.poller = None
         self.distinct_ids_feature_flags_reported = SizeLimitedDict(MAX_DICT_SIZE, set)
         self.disabled = disabled
@@ -302,7 +311,9 @@ class Client(object):
             self.log.setLevel(logging.WARNING)
 
         if self.enable_exception_autocapture:
-            self.exception_capture = ExceptionCapture(self, integrations=self.exception_autocapture_integrations)
+            self.exception_capture = ExceptionCapture(
+                self, integrations=self.exception_autocapture_integrations
+            )
 
         if sync_mode:
             self.consumers = None
@@ -348,12 +359,24 @@ class Client(object):
         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"
+        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):
+    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.",
@@ -376,34 +399,60 @@ class Client(object):
         return self._enqueue(msg, disable_geoip)
 
     def get_feature_variants(
-        self, distinct_id, groups=None, person_properties=None, group_properties=None, disable_geoip=None
+        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_flags_decision(distinct_id, groups, person_properties, group_properties, disable_geoip)
+        resp_data = self.get_flags_decision(
+            distinct_id, groups, person_properties, group_properties, disable_geoip
+        )
         return to_values(resp_data) or {}
 
     def get_feature_payloads(
-        self, distinct_id, groups=None, person_properties=None, group_properties=None, disable_geoip=None
+        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_flags_decision(distinct_id, groups, person_properties, group_properties, disable_geoip)
+        resp_data = self.get_flags_decision(
+            distinct_id, groups, person_properties, group_properties, disable_geoip
+        )
         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
+        self,
+        distinct_id,
+        groups=None,
+        person_properties=None,
+        group_properties=None,
+        disable_geoip=None,
     ) -> FlagsAndPayloads:
         """
         Get feature flags and payloads for a distinct_id by calling decide.
         """
-        resp = self.get_flags_decision(distinct_id, groups, person_properties, group_properties, disable_geoip)
+        resp = self.get_flags_decision(
+            distinct_id, groups, person_properties, group_properties, disable_geoip
+        )
         return to_flags_and_payloads(resp)
 
     def get_flags_decision(
-        self, distinct_id, groups=None, person_properties=None, group_properties=None, disable_geoip=None
+        self,
+        distinct_id,
+        groups=None,
+        person_properties=None,
+        group_properties=None,
+        disable_geoip=None,
     ) -> FlagsResponse:
         """
         Get feature flags decision, using either flags() or decide() API based on rollout.
@@ -423,20 +472,29 @@ class Client(object):
             "groups": groups,
             "person_properties": person_properties,
             "group_properties": group_properties,
-            "disable_geoip": disable_geoip,
+            "geoip_disable": disable_geoip,
         }
 
         use_flags = is_token_in_rollout(
-            self.api_key, ROLLOUT_PERCENTAGE, included_hashes=INCLUDED_HASHES, excluded_hashes=EXCLUDED_HASHES
+            self.api_key,
+            ROLLOUT_PERCENTAGE,
+            included_hashes=INCLUDED_HASHES,
+            excluded_hashes=EXCLUDED_HASHES,
         )
 
         if use_flags:
             resp_data = flags(
-                self.api_key, self.host, timeout=self.feature_flags_request_timeout_seconds, **request_data
+                self.api_key,
+                self.host,
+                timeout=self.feature_flags_request_timeout_seconds,
+                **request_data,
             )
         else:
             resp_data = decide(
-                self.api_key, self.host, timeout=self.feature_flags_request_timeout_seconds, **request_data
+                self.api_key,
+                self.host,
+                timeout=self.feature_flags_request_timeout_seconds,
+                **request_data,
             )
 
         return normalize_flags_response(resp_data)
@@ -466,6 +524,11 @@ class Client(object):
         require("properties", properties, dict)
         require("event", event, string_types)
 
+        # Grab current context tags, if any exist
+        context_tags = get_tags()
+        if context_tags:
+            properties.update(context_tags)
+
         msg = {
             "properties": properties,
             "timestamp": timestamp,
@@ -482,20 +545,31 @@ class Client(object):
         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)
+                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}")
+                self.log.exception(
+                    f"[FEATURE FLAGS] Unable to get feature variants: {e}"
+                )
 
         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
+                distinct_id,
+                groups=(groups or {}),
+                disable_geoip=disable_geoip,
+                only_evaluate_locally=True,
             )
 
         for feature, variant in (feature_variants or {}).items():
             extra_properties[f"$feature/{feature}"] = variant
 
-        active_feature_flags = [key for (key, value) in (feature_variants or {}).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
 
@@ -504,7 +578,15 @@ class Client(object):
 
         return self._enqueue(msg, disable_geoip)
 
-    def set(self, distinct_id=None, properties=None, context=None, timestamp=None, uuid=None, disable_geoip=None):
+    def set(
+        self,
+        distinct_id=None,
+        properties=None,
+        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.",
@@ -526,7 +608,15 @@ class Client(object):
 
         return self._enqueue(msg, disable_geoip)
 
-    def set_once(self, distinct_id=None, properties=None, context=None, timestamp=None, uuid=None, disable_geoip=None):
+    def set_once(
+        self,
+        distinct_id=None,
+        properties=None,
+        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.",
@@ -589,7 +679,15 @@ class Client(object):
 
         return self._enqueue(msg, disable_geoip)
 
-    def alias(self, previous_id=None, distinct_id=None, context=None, timestamp=None, uuid=None, disable_geoip=None):
+    def alias(
+        self,
+        previous_id=None,
+        distinct_id=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.",
@@ -613,7 +711,14 @@ class Client(object):
         return self._enqueue(msg, disable_geoip)
 
     def page(
-        self, distinct_id=None, url=None, properties=None, context=None, timestamp=None, uuid=None, disable_geoip=None
+        self,
+        distinct_id=None,
+        url=None,
+        properties=None,
+        context=None,
+        timestamp=None,
+        uuid=None,
+        disable_geoip=None,
     ):
         if context is not None:
             warnings.warn(
@@ -662,6 +767,13 @@ class Client(object):
         try:
             properties = properties or {}
 
+            # Check if this exception has already been captured
+            if exception is not None and hasattr(
+                exception, "__posthog_exception_captured"
+            ):
+                self.log.debug("Exception already captured, skipping")
+                return
+
             # if there's no distinct_id, we'll generate one and set personless mode
             # via $process_person_profile = false
             if distinct_id is None:
@@ -696,7 +808,9 @@ class Client(object):
 
             properties = {
                 "$exception_type": all_exceptions_with_trace_and_in_app[0].get("type"),
-                "$exception_message": all_exceptions_with_trace_and_in_app[0].get("value"),
+                "$exception_message": all_exceptions_with_trace_and_in_app[0].get(
+                    "value"
+                ),
                 "$exception_list": all_exceptions_with_trace_and_in_app,
                 "$exception_personURL": f"{remove_trailing_slash(self.raw_host)}/project/{self.api_key}/person/{distinct_id}",
                 **properties,
@@ -705,7 +819,15 @@ class Client(object):
             if self.log_captured_exceptions:
                 self.log.exception(exception, extra=kwargs)
 
-            return self.capture(distinct_id, "$exception", properties, context, timestamp, uuid, groups)
+            res = self.capture(
+                distinct_id, "$exception", properties, context, timestamp, uuid, groups
+            )
+
+            # Mark the exception as captured to prevent duplicate captures
+            if exception is not None:
+                setattr(exception, "__posthog_exception_captured", True)
+
+            return res
         except Exception as e:
             self.log.exception(f"Failed to capture exception: {e}")
 
@@ -858,13 +980,18 @@ class Client(object):
 
     def load_feature_flags(self):
         if not self.personal_api_key:
-            self.log.warning("[FEATURE FLAGS] You have to specify a personal_api_key to use feature flags.")
+            self.log.warning(
+                "[FEATURE FLAGS] You have to specify a personal_api_key to use feature flags."
+            )
             self.feature_flags = []
             return
 
         self._load_feature_flags()
         if not (self.poller and self.poller.is_alive()):
-            self.poller = Poller(interval=timedelta(seconds=self.poll_interval), execute=self._load_feature_flags)
+            self.poller = Poller(
+                interval=timedelta(seconds=self.poll_interval),
+                execute=self._load_feature_flags,
+            )
             self.poller.start()
 
     def _compute_flag_locally(
@@ -909,9 +1036,13 @@ class Client(object):
                 return False
 
             focused_group_properties = group_properties[group_name]
-            return match_feature_flag_properties(feature_flag, groups[group_name], focused_group_properties)
+            return match_feature_flag_properties(
+                feature_flag, groups[group_name], focused_group_properties
+            )
         else:
-            return match_feature_flag_properties(feature_flag, distinct_id, person_properties, self.cohorts)
+            return match_feature_flag_properties(
+                feature_flag, distinct_id, person_properties, self.cohorts
+            )
 
     def feature_enabled(
         self,
@@ -960,28 +1091,47 @@ class Client(object):
         if self.disabled:
             return None
 
-        person_properties, group_properties = self._add_local_person_and_group_properties(
-            distinct_id, groups, person_properties, group_properties
+        person_properties, group_properties = (
+            self._add_local_person_and_group_properties(
+                distinct_id, groups, person_properties, group_properties
+            )
         )
 
         flag_result = None
         flag_details = None
         request_id = None
 
-        flag_value = self._locally_evaluate_flag(key, distinct_id, groups, person_properties, group_properties)
+        flag_value = self._locally_evaluate_flag(
+            key, distinct_id, groups, person_properties, group_properties
+        )
         flag_was_locally_evaluated = flag_value is not None
 
         if flag_was_locally_evaluated:
             lookup_match_value = override_match_value or flag_value
-            payload = self._compute_payload_locally(key, lookup_match_value) if lookup_match_value else None
-            flag_result = FeatureFlagResult.from_value_and_payload(key, lookup_match_value, payload)
+            payload = (
+                self._compute_payload_locally(key, lookup_match_value)
+                if lookup_match_value
+                else None
+            )
+            flag_result = FeatureFlagResult.from_value_and_payload(
+                key, lookup_match_value, payload
+            )
         elif not only_evaluate_locally:
             try:
                 flag_details, request_id = self._get_feature_flag_details_from_decide(
-                    key, distinct_id, groups, person_properties, group_properties, disable_geoip
+                    key,
+                    distinct_id,
+                    groups,
+                    person_properties,
+                    group_properties,
+                    disable_geoip,
+                )
+                flag_result = FeatureFlagResult.from_flag_details(
+                    flag_details, override_match_value
+                )
+                self.log.debug(
+                    f"Successfully computed flag remotely: #{key} -> #{flag_result}"
                 )
-                flag_result = FeatureFlagResult.from_flag_details(flag_details, override_match_value)
-                self.log.debug(f"Successfully computed flag remotely: #{key} -> #{flag_result}")
             except Exception as e:
                 self.log.exception(f"[FEATURE FLAGS] Unable to get flag remotely: {e}")
 
@@ -1073,9 +1223,9 @@ class Client(object):
         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"
+            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:
@@ -1087,11 +1237,15 @@ class Client(object):
                         person_properties=person_properties,
                         group_properties=group_properties,
                     )
-                    self.log.debug(f"Successfully computed flag locally: {key} -> {response}")
+                    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}")
+                    self.log.exception(
+                        f"[FEATURE FLAGS] Error while computing variant locally: {e}"
+                    )
         return response
 
     def get_feature_flag_payload(
@@ -1132,7 +1286,9 @@ class Client(object):
         """
         Calls /decide and returns the flag details and request id
         """
-        resp_data = self.get_flags_decision(distinct_id, groups, person_properties, group_properties, disable_geoip)
+        resp_data = self.get_flags_decision(
+            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
@@ -1150,9 +1306,14 @@ class Client(object):
         request_id: Optional[str],
         flag_details: Optional[FeatureFlag],
     ):
-        feature_flag_reported_key = f"{key}_{'::null::' if response is None else str(response)}"
+        feature_flag_reported_key = (
+            f"{key}_{'::null::' if response is None else str(response)}"
+        )
 
-        if feature_flag_reported_key not in self.distinct_ids_feature_flags_reported[distinct_id]:
+        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,
@@ -1171,7 +1332,9 @@ class Client(object):
                     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
+                        properties["$feature_flag_version"] = (
+                            flag_details.metadata.version
+                        )
                     if flag_details.metadata.id:
                         properties["$feature_flag_id"] = flag_details.metadata.id
 
@@ -1182,7 +1345,9 @@ class Client(object):
                 groups=groups,
                 disable_geoip=disable_geoip,
             )
-            self.distinct_ids_feature_flags_reported[distinct_id].add(feature_flag_reported_key)
+            self.distinct_ids_feature_flags_reported[distinct_id].add(
+                feature_flag_reported_key
+            )
 
     def get_remote_config_payload(self, key: str):
         if self.disabled:
@@ -1202,9 +1367,13 @@ class Client(object):
                 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}")
+            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]:
+    def _compute_payload_locally(
+        self, key: str, match_value: FlagValue
+    ) -> Optional[str]:
         payload = None
 
         if self.feature_flags_by_key is None:
@@ -1216,7 +1385,11 @@ class Client(object):
             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)
+            lookup_value = (
+                "true"
+                if isinstance(match_value, bool) and match_value
+                else str(match_value)
+            )
             payload = flag_payloads.get(lookup_value, None)
         return payload
 
@@ -1254,12 +1427,17 @@ class Client(object):
         if self.disabled:
             return {"featureFlags": None, "featureFlagPayloads": None}
 
-        person_properties, group_properties = self._add_local_person_and_group_properties(
-            distinct_id, groups, person_properties, group_properties
+        person_properties, group_properties = (
+            self._add_local_person_and_group_properties(
+                distinct_id, groups, person_properties, group_properties
+            )
         )
 
         response, fallback_to_decide = self._get_all_flags_and_payloads_locally(
-            distinct_id, groups=groups, person_properties=person_properties, group_properties=group_properties
+            distinct_id,
+            groups=groups,
+            person_properties=person_properties,
+            group_properties=group_properties,
         )
 
         if fallback_to_decide and not only_evaluate_locally:
@@ -1273,12 +1451,20 @@ class Client(object):
                 )
                 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}")
+                self.log.exception(
+                    f"[FEATURE FLAGS] Unable to get feature flags and payloads: {e}"
+                )
 
         return response
 
     def _get_all_flags_and_payloads_locally(
-        self, distinct_id, *, groups={}, person_properties={}, group_properties={}, warn_on_unknown_groups=False
+        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)
@@ -1301,25 +1487,37 @@ class Client(object):
                         group_properties=group_properties,
                         warn_on_unknown_groups=warn_on_unknown_groups,
                     )
-                    matched_payload = self._compute_payload_locally(flag["key"], flags[flag["key"]])
+                    matched_payload = self._compute_payload_locally(
+                        flag["key"], flags[flag["key"]]
+                    )
                     if matched_payload:
                         payloads[flag["key"]] = matched_payload
                 except InconclusiveMatchError:
                     # No need to log this, since it's just telling us to fall back to `/decide`
                     fallback_to_decide = True
                 except Exception as e:
-                    self.log.exception(f"[FEATURE FLAGS] Error while computing variant and payload: {e}")
+                    self.log.exception(
+                        f"[FEATURE FLAGS] Error while computing variant and payload: {e}"
+                    )
                     fallback_to_decide = True
         else:
             fallback_to_decide = True
 
-        return {"featureFlags": flags, "featureFlagPayloads": payloads}, fallback_to_decide
+        return {
+            "featureFlags": flags,
+            "featureFlagPayloads": payloads,
+        }, fallback_to_decide
 
     def feature_flag_definitions(self):
         return self.feature_flags
 
-    def _add_local_person_and_group_properties(self, distinct_id, groups, person_properties, group_properties):
-        all_person_properties = {"distinct_id": distinct_id, **(person_properties or {})}
+    def _add_local_person_and_group_properties(
+        self, distinct_id, groups, person_properties, group_properties
+    ):
+        all_person_properties = {
+            "distinct_id": distinct_id,
+            **(person_properties or {}),
+        }
 
         all_group_properties = {}
         if groups:

posthog/consumer.py

@@ -107,7 +107,9 @@ class Consumer(Thread):
                 item = queue.get(block=True, timeout=self.flush_interval - elapsed)
                 item_size = len(json.dumps(item, cls=DatetimeSerializer).encode())
                 if item_size > MAX_MSG_SIZE:
-                    self.log.error("Item exceeds 900kib limit, dropping. (%s)", str(item))
+                    self.log.error(
+                        "Item exceeds 900kib limit, dropping. (%s)", str(item)
+                    )
                     continue
                 items.append(item)
                 total_size += item_size
@@ -134,7 +136,9 @@ class Consumer(Thread):
                 # retry on all other errors (eg. network)
                 return False
 
-        @backoff.on_exception(backoff.expo, Exception, max_tries=self.retries + 1, giveup=fatal_exception)
+        @backoff.on_exception(
+            backoff.expo, Exception, max_tries=self.retries + 1, giveup=fatal_exception
+        )
         def send_request():
             batch_post(
                 self.api_key,

posthog/exception_capture.py

@@ -22,7 +22,9 @@ class ExceptionCapture:
 
     log = logging.getLogger("posthog")
 
-    def __init__(self, client: "Client", integrations: Optional[List[Integrations]] = None):
+    def __init__(
+        self, client: "Client", integrations: Optional[List[Integrations]] = None
+    ):
         self.client = client
         self.original_excepthook = sys.excepthook
         sys.excepthook = self.exception_handler

posthog/exception_integrations/django.py

@@ -32,7 +32,6 @@ class DjangoIntegration:
     identifier = "django"
 
     def __init__(self, capture_exception_fn=None):
-
         if DJANGO_VERSION < (4, 2):
             raise IntegrationEnablingError("Django 4.2 or newer is required.")
 
@@ -60,7 +59,6 @@ class DjangoIntegration:
 
 
 class DjangoRequestExtractor:
-
     def __init__(self, request):
         # type: (Any) -> None
         self.request = request

posthog/exception_utils.py

@@ -24,7 +24,6 @@ DEFAULT_MAX_VALUE_LENGTH = 1024
 
 
 if TYPE_CHECKING:
-
     from types import FrameType, TracebackType
     from typing import (  # noqa: F401
         Any,
@@ -52,7 +51,9 @@ if TYPE_CHECKING:
     Event = TypedDict(
         "Event",
         {
-            "breadcrumbs": Dict[Literal["values"], List[Dict[str, Any]]],  # TODO: We can expand on this type
+            "breadcrumbs": Dict[
+                Literal["values"], List[Dict[str, Any]]
+            ],  # TODO: We can expand on this type
             "check_in_id": str,
             "contexts": Dict[str, Dict[str, object]],
             "dist": str,
@@ -60,7 +61,9 @@ if TYPE_CHECKING:
             "environment": str,
             "errors": List[Dict[str, Any]],  # TODO: We can expand on this type
             "event_id": str,
-            "exception": Dict[Literal["values"], List[Dict[str, Any]]],  # TODO: We can expand on this type
+            "exception": Dict[
+                Literal["values"], List[Dict[str, Any]]
+            ],  # TODO: We can expand on this type
             # "extra": MutableMapping[str, object],
             # "fingerprint": List[str],
             "level": LogLevelStr,
@@ -78,13 +81,17 @@ if TYPE_CHECKING:
             # "sdk": Mapping[str, object],
             "server_name": str,
             "spans": List[Dict[str, object]],
-            "stacktrace": Dict[str, object],  # We access this key in the code, but I am unsure whether we ever set it
+            "stacktrace": Dict[
+                str, object
+            ],  # We access this key in the code, but I am unsure whether we ever set it
             "start_timestamp": datetime,
             "status": Optional[str],
             # "tags": MutableMapping[
             #     str, str
             # ],  # Tags must be less than 200 characters each
-            "threads": Dict[Literal["values"], List[Dict[str, Any]]],  # TODO: We can expand on this type
+            "threads": Dict[
+                Literal["values"], List[Dict[str, Any]]
+            ],  # TODO: We can expand on this type
             "timestamp": Optional[datetime],  # Must be set before sending the event
             "transaction": str,
             # "transaction_info": Mapping[str, Any],  # TODO: We can expand on this type
@@ -273,7 +280,10 @@ def get_lines_from_file(
     upper_bound = min(lineno + 1 + context_lines, len(source))
 
     try:
-        pre_context = [strip_string(line.strip("\r\n"), max_length=max_length) for line in source[lower_bound:lineno]]
+        pre_context = [
+            strip_string(line.strip("\r\n"), max_length=max_length)
+            for line in source[lower_bound:lineno]
+        ]
         context_line = strip_string(source[lineno].strip("\r\n"), max_length=max_length)
         post_context = [
             strip_string(line.strip("\r\n"), max_length=max_length)
@@ -305,7 +315,9 @@ def get_source_context(
         loader = None
     lineno = tb_lineno - 1
     if lineno is not None and abs_path:
-        return get_lines_from_file(abs_path, lineno, max_value_length, loader=loader, module=module)
+        return get_lines_from_file(
+            abs_path, lineno, max_value_length, loader=loader, module=module
+        )
     return [], None, []
 
 
@@ -342,7 +354,9 @@ def filename_for_module(module, abs_path):
         if not base_module_path:
             return abs_path
 
-        return abs_path.split(base_module_path.rsplit(os.sep, 2)[0], 1)[-1].lstrip(os.sep)
+        return abs_path.split(base_module_path.rsplit(os.sep, 2)[0], 1)[-1].lstrip(
+            os.sep
+        )
     except Exception:
         return abs_path
 
@@ -431,7 +445,11 @@ def get_errno(exc_value):
 
 def get_error_message(exc_value):
     # type: (Optional[BaseException]) -> str
-    return getattr(exc_value, "message", "") or getattr(exc_value, "detail", "") or safe_str(exc_value)
+    return (
+        getattr(exc_value, "message", "")
+        or getattr(exc_value, "detail", "")
+        or safe_str(exc_value)
+    )
 
 
 def single_exception_from_error_tuple(
@@ -452,7 +470,9 @@ def single_exception_from_error_tuple(
     https://develop.sentry.dev/sdk/event-payloads/exception/
     """
     exception_value = {}  # type: Dict[str, Any]
-    exception_value["mechanism"] = mechanism.copy() if mechanism else {"type": "generic", "handled": True}
+    exception_value["mechanism"] = (
+        mechanism.copy() if mechanism else {"type": "generic", "handled": True}
+    )
     if exception_id is not None:
         exception_value["mechanism"]["exception_id"] = exception_id
 
@@ -462,7 +482,9 @@ def single_exception_from_error_tuple(
         errno = None
 
     if errno is not None:
-        exception_value["mechanism"].setdefault("meta", {}).setdefault("errno", {}).setdefault("number", errno)
+        exception_value["mechanism"].setdefault("meta", {}).setdefault(
+            "errno", {}
+        ).setdefault("number", errno)
 
     if source is not None:
         exception_value["mechanism"]["source"] = source
@@ -475,7 +497,9 @@ def single_exception_from_error_tuple(
     if is_root_exception and "type" not in exception_value["mechanism"]:
         exception_value["mechanism"]["type"] = "generic"
 
-    is_exception_group = BaseExceptionGroup is not None and isinstance(exc_value, BaseExceptionGroup)
+    is_exception_group = BaseExceptionGroup is not None and isinstance(
+        exc_value, BaseExceptionGroup
+    )
     if is_exception_group:
         exception_value["mechanism"]["is_exception_group"] = True
 
@@ -523,7 +547,11 @@ if HAS_CHAINED_EXCEPTIONS:
         seen_exceptions = []
         seen_exception_ids = set()  # type: Set[int]
 
-        while exc_type is not None and exc_value is not None and id(exc_value) not in seen_exception_ids:
+        while (
+            exc_type is not None
+            and exc_value is not None
+            and id(exc_value) not in seen_exception_ids
+        ):
             yield exc_type, exc_value, tb
 
             # Avoid hashing random types we don't know anything
@@ -583,11 +611,17 @@ def exceptions_from_error(
     parent_id = exception_id
     exception_id += 1
 
-    should_supress_context = hasattr(exc_value, "__suppress_context__") and exc_value.__suppress_context__  # type: ignore
+    should_supress_context = (
+        hasattr(exc_value, "__suppress_context__") and exc_value.__suppress_context__  # type: ignore
+    )
     if should_supress_context:
         # Add direct cause.
         # The field `__cause__` is set when raised with the exception (using the `from` keyword).
-        exception_has_cause = exc_value and hasattr(exc_value, "__cause__") and exc_value.__cause__ is not None
+        exception_has_cause = (
+            exc_value
+            and hasattr(exc_value, "__cause__")
+            and exc_value.__cause__ is not None
+        )
         if exception_has_cause:
             cause = exc_value.__cause__  # type: ignore
             (exception_id, child_exceptions) = exceptions_from_error(
@@ -604,7 +638,11 @@ def exceptions_from_error(
     else:
         # Add indirect cause.
         # The field `__context__` is assigned if another exception occurs while handling the exception.
-        exception_has_content = exc_value and hasattr(exc_value, "__context__") and exc_value.__context__ is not None
+        exception_has_content = (
+            exc_value
+            and hasattr(exc_value, "__context__")
+            and exc_value.__context__ is not None
+        )
         if exception_has_content:
             context = exc_value.__context__  # type: ignore
             (exception_id, child_exceptions) = exceptions_from_error(
@@ -645,7 +683,9 @@ def exceptions_from_error_tuple(
     # type: (...) -> List[Dict[str, Any]]
     exc_type, exc_value, tb = exc_info
 
-    is_exception_group = BaseExceptionGroup is not None and isinstance(exc_value, BaseExceptionGroup)
+    is_exception_group = BaseExceptionGroup is not None and isinstance(
+        exc_value, BaseExceptionGroup
+    )
 
     if is_exception_group:
         (_, exceptions) = exceptions_from_error(
@@ -661,7 +701,11 @@ def exceptions_from_error_tuple(
     else:
         exceptions = []
         for exc_type, exc_value, tb in walk_exception_chain(exc_info):
-            exceptions.append(single_exception_from_error_tuple(exc_type, exc_value, tb, client_options, mechanism))
+            exceptions.append(
+                single_exception_from_error_tuple(
+                    exc_type, exc_value, tb, client_options, mechanism
+                )
+            )
 
     exceptions.reverse()
 
@@ -789,7 +833,11 @@ def event_from_exception(
     return (
         {
             "level": "error",
-            "exception": {"values": exceptions_from_error_tuple(exc_info, client_options, mechanism)},
+            "exception": {
+                "values": exceptions_from_error_tuple(
+                    exc_info, client_options, mechanism
+                )
+            },
         },
         hint,
     )
@@ -813,7 +861,9 @@ def _module_in_list(name, items):
 def _is_external_source(abs_path):
     # type: (str) -> bool
     # check if frame is in 'site-packages' or 'dist-packages'
-    external_source = re.search(r"[\\/](?:dist|site)-packages[\\/]", abs_path) is not None
+    external_source = (
+        re.search(r"[\\/](?:dist|site)-packages[\\/]", abs_path) is not None
+    )
     return external_source
 
 

posthog/feature_flags.py

@@ -43,20 +43,28 @@ def get_matching_variant(flag, distinct_id):
 def variant_lookup_table(feature_flag):
     lookup_table = []
     value_min = 0
-    multivariates = ((feature_flag.get("filters") or {}).get("multivariate") or {}).get("variants") or []
+    multivariates = ((feature_flag.get("filters") or {}).get("multivariate") or {}).get(
+        "variants"
+    ) or []
     for variant in multivariates:
         value_max = value_min + variant["rollout_percentage"] / 100
-        lookup_table.append({"value_min": value_min, "value_max": value_max, "key": variant["key"]})
+        lookup_table.append(
+            {"value_min": value_min, "value_max": value_max, "key": variant["key"]}
+        )
         value_min = value_max
     return lookup_table
 
 
-def match_feature_flag_properties(flag, distinct_id, properties, cohort_properties=None) -> FlagValue:
+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 []
+    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
@@ -70,7 +78,9 @@ def match_feature_flag_properties(flag, distinct_id, properties, cohort_properti
         try:
             # if any one condition resolves to True, we can shortcircuit and return
             # the matching variant
-            if is_condition_match(flag, distinct_id, condition, properties, cohort_properties):
+            if is_condition_match(
+                flag, distinct_id, condition, properties, cohort_properties
+            ):
                 variant_override = condition.get("variant")
                 if variant_override and variant_override in valid_variant_keys:
                     variant = variant_override
@@ -81,14 +91,18 @@ def match_feature_flag_properties(flag, distinct_id, properties, cohort_properti
             is_inconclusive = True
 
     if is_inconclusive:
-        raise InconclusiveMatchError("Can't determine if feature flag is enabled or not with given properties")
+        raise InconclusiveMatchError(
+            "Can't determine if feature flag is enabled or not with given properties"
+        )
 
     # We can only return False when either all conditions are False, or
     # no condition was inconclusive.
     return False
 
 
-def is_condition_match(feature_flag, distinct_id, condition, properties, cohort_properties) -> bool:
+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"):
@@ -103,7 +117,9 @@ def is_condition_match(feature_flag, distinct_id, condition, properties, cohort_
         if rollout_percentage is None:
             return True
 
-    if rollout_percentage is not None and _hash(feature_flag["key"], distinct_id) > (rollout_percentage / 100):
+    if rollout_percentage is not None and _hash(feature_flag["key"], distinct_id) > (
+        rollout_percentage / 100
+    ):
         return False
 
     return True
@@ -117,7 +133,9 @@ def match_property(property, property_values) -> bool:
     value = property.get("value")
 
     if key not in property_values:
-        raise InconclusiveMatchError("can't match properties without a given property value")
+        raise InconclusiveMatchError(
+            "can't match properties without a given property value"
+        )
 
     if operator == "is_not_set":
         raise InconclusiveMatchError("can't match properties with operator is_not_set")
@@ -131,7 +149,9 @@ def match_property(property, property_values) -> bool:
 
         def compute_exact_match(value, override_value):
             if isinstance(value, list):
-                return str(override_value).casefold() in [str(val).casefold() for val in value]
+                return str(override_value).casefold() in [
+                    str(val).casefold() for val in value
+                ]
             return utils.str_iequals(value, override_value)
 
         if operator == "exact":
@@ -149,10 +169,16 @@ def match_property(property, property_values) -> bool:
         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
+        return (
+            is_valid_regex(str(value))
+            and re.compile(str(value)).search(str(override_value)) is not None
+        )
 
     if operator == "not_regex":
-        return is_valid_regex(str(value)) and re.compile(str(value)).search(str(override_value)) is None
+        return (
+            is_valid_regex(str(value))
+            and re.compile(str(value)).search(str(override_value)) is None
+        )
 
     if operator in ("gt", "gte", "lt", "lte"):
         # :TRICKY: We adjust comparison based on the override value passed in,
@@ -191,10 +217,14 @@ def match_property(property, property_values) -> bool:
                 parsed_date = parser.parse(str(value))
                 parsed_date = convert_to_datetime_aware(parsed_date)
         except Exception as e:
-            raise InconclusiveMatchError("The date set on the flag is not a valid format") from e
+            raise InconclusiveMatchError(
+                "The date set on the flag is not a valid format"
+            ) from e
 
         if not parsed_date:
-            raise InconclusiveMatchError("The date set on the flag is not a valid format")
+            raise InconclusiveMatchError(
+                "The date set on the flag is not a valid format"
+            )
 
         if isinstance(override_value, datetime.datetime):
             override_date = convert_to_datetime_aware(override_value)
@@ -218,7 +248,9 @@ def match_property(property, property_values) -> bool:
             except Exception:
                 raise InconclusiveMatchError("The date provided is not a valid format")
         else:
-            raise InconclusiveMatchError("The date provided must be a string or date object")
+            raise InconclusiveMatchError(
+                "The date provided must be a string or date object"
+            )
 
     # if we get here, we don't know how to handle the operator
     raise InconclusiveMatchError(f"Unknown operator {operator}")
@@ -236,7 +268,9 @@ def match_cohort(property, property_values, cohort_properties) -> bool:
     # }
     cohort_id = str(property.get("value"))
     if cohort_id not in cohort_properties:
-        raise InconclusiveMatchError("can't match cohort without a given cohort property value")
+        raise InconclusiveMatchError(
+            "can't match cohort without a given cohort property value"
+        )
 
     property_group = cohort_properties[cohort_id]
     return match_property_group(property_group, property_values, cohort_properties)
@@ -272,7 +306,9 @@ def match_property_group(property_group, property_values, cohort_properties) ->
                 error_matching_locally = True
 
         if error_matching_locally:
-            raise InconclusiveMatchError("Can't match cohort without a given cohort property value")
+            raise InconclusiveMatchError(
+                "Can't match cohort without a given cohort property value"
+            )
         # if we get here, all matched in AND case, or none matched in OR case
         return property_group_type == "AND"
 
@@ -303,13 +339,17 @@ def match_property_group(property_group, property_values, cohort_properties) ->
                 error_matching_locally = True
 
         if error_matching_locally:
-            raise InconclusiveMatchError("can't match cohort without a given cohort property value")
+            raise InconclusiveMatchError(
+                "can't match cohort without a given cohort property value"
+            )
 
         # if we get here, all matched in AND case, or none matched in OR case
         return property_group_type == "AND"
 
 
-def relative_date_parse_for_feature_flag_matching(value: str) -> Optional[datetime.datetime]:
+def relative_date_parse_for_feature_flag_matching(
+    value: str,
+) -> Optional[datetime.datetime]:
     regex = r"^-?(?P<number>[0-9]+)(?P<interval>[a-z])$"
     match = re.search(regex, value)
     parsed_dt = datetime.datetime.now(datetime.timezone.utc)

posthog/request.py

@@ -43,7 +43,12 @@ def determine_server_host(host: Optional[str]) -> str:
 
 
 def post(
-    api_key: str, host: Optional[str] = None, path=None, gzip: bool = False, timeout: int = 15, **kwargs
+    api_key: str,
+    host: Optional[str] = None,
+    path=None,
+    gzip: bool = False,
+    timeout: int = 15,
+    **kwargs,
 ) -> requests.Response:
     """Post the `kwargs` to the API"""
     log = logging.getLogger("posthog")
@@ -100,34 +105,67 @@ def _process_response(
         raise APIError(res.status_code, res.text)
 
 
-def decide(api_key: str, host: Optional[str] = None, gzip: bool = False, timeout: int = 15, **kwargs) -> Any:
+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=4", gzip, timeout, **kwargs)
     return _process_response(res, success_message="Feature flags decided successfully")
 
 
-def flags(api_key: str, host: Optional[str] = None, gzip: bool = False, timeout: int = 15, **kwargs) -> Any:
+def flags(
+    api_key: str,
+    host: Optional[str] = None,
+    gzip: bool = False,
+    timeout: int = 15,
+    **kwargs,
+) -> Any:
     """Post the `kwargs to the flags API endpoint"""
     res = post(api_key, host, "/flags/?v=2", gzip, timeout, **kwargs)
-    return _process_response(res, success_message="Feature flags evaluated successfully")
+    return _process_response(
+        res, success_message="Feature flags evaluated successfully"
+    )
 
 
-def remote_config(personal_api_key: str, host: Optional[str] = None, key: str = "", timeout: int = 15) -> Any:
+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)
+    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
+    api_key: str,
+    host: Optional[str] = None,
+    gzip: bool = False,
+    timeout: int = 15,
+    **kwargs,
 ) -> requests.Response:
     """Post the `kwargs` to the batch API endpoint for events"""
     res = post(api_key, host, "/batch/", gzip, timeout, **kwargs)
-    return _process_response(res, success_message="data uploaded successfully", return_json=False)
+    return _process_response(
+        res, success_message="data uploaded successfully", return_json=False
+    )
 
 
-def get(api_key: str, url: str, host: Optional[str] = None, timeout: Optional[int] = None) -> requests.Response:
+def get(
+    api_key: str, url: str, host: Optional[str] = None, timeout: Optional[int] = None
+) -> requests.Response:
     url = remove_trailing_slash(host or DEFAULT_HOST) + url
-    res = requests.get(url, headers={"Authorization": "Bearer %s" % api_key, "User-Agent": USER_AGENT}, timeout=timeout)
+    res = requests.get(
+        url,
+        headers={"Authorization": "Bearer %s" % api_key, "User-Agent": USER_AGENT},
+        timeout=timeout,
+    )
     return _process_response(res, success_message=f"GET {url} completed successfully")
 
 

posthog/scopes.py

@@ -0,0 +1,122 @@
+import contextvars
+from contextlib import contextmanager
+from typing import Any, Callable, Dict, TypeVar, cast
+
+_context_stack: contextvars.ContextVar[list] = contextvars.ContextVar(
+    "posthog_context_stack", default=[{}]
+)
+
+
+def _get_current_context() -> Dict[str, Any]:
+    return _context_stack.get()[-1]
+
+
+@contextmanager
+def new_context(fresh=False):
+    """
+     Create a new context scope that will be active for the duration of the with block.
+     Any tags set within this scope will be isolated to this context. Any exceptions raised
+    or events captured within the context will be tagged with the context tags.
+
+     Args:
+         fresh: Whether to start with a fresh context (default: False).
+                If False, inherits tags from parent context.
+                If True, starts with no tags.
+
+     Examples:
+         # Inherit parent context tags
+         with posthog.new_context():
+             posthog.tag("request_id", "123")
+             # Both this event and the exception will be tagged with the context tags
+             posthog.capture("event_name", {"property": "value"})
+             raise ValueError("Something went wrong")
+
+         # Start with fresh context (no inherited tags)
+         with posthog.new_context(fresh=True):
+             posthog.tag("request_id", "123")
+             # Both this event and the exception will be tagged with the context tags
+             posthog.capture("event_name", {"property": "value"})
+             raise ValueError("Something went wrong")
+
+    """
+    import posthog
+
+    current_tags = _get_current_context().copy()
+    current_stack = _context_stack.get()
+    new_stack = current_stack + [{}] if fresh else current_stack + [current_tags]
+    token = _context_stack.set(new_stack)
+
+    try:
+        yield
+    except Exception as e:
+        posthog.capture_exception(e)
+        raise
+    finally:
+        _context_stack.reset(token)
+
+
+def tag(key: str, value: Any) -> None:
+    """
+    Add a tag to the current context.
+
+    Args:
+        key: The tag key
+        value: The tag value
+
+    Example:
+        posthog.tag("user_id", "123")
+    """
+    _get_current_context()[key] = value
+
+
+def get_tags() -> Dict[str, Any]:
+    """
+    Get all tags from the current context. Note, modifying
+    the returned dictionary will not affect the current context.
+
+    Returns:
+        Dict of all tags in the current context
+    """
+    return _get_current_context().copy()
+
+
+def clear_tags() -> None:
+    """Clear all tags in the current context."""
+    _get_current_context().clear()
+
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def scoped(fresh=False):
+    """
+    Decorator that creates a new context for the function. Simply wraps
+    the function in a with posthog.new_context(): block.
+
+    Args:
+        fresh: Whether to start with a fresh context (default: False)
+
+    Example:
+        @posthog.scoped()
+        def process_payment(payment_id):
+            posthog.tag("payment_id", payment_id)
+            posthog.tag("payment_method", "credit_card")
+
+            # This event will be captured with tags
+            posthog.capture("payment_started")
+            # If this raises an exception, it will be captured with tags
+            # and then re-raised
+            some_risky_function()
+    """
+
+    def decorator(func: F) -> F:
+        from functools import wraps
+
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            with new_context(fresh=fresh):
+                return func(*args, **kwargs)
+
+        return cast(F, wrapper)
+
+    return decorator

posthog/sentry/posthog_integration.py

@@ -17,7 +17,9 @@ if MYPY:
 class PostHogIntegration(Integration):
     identifier = "posthog-python"
     organization = None  # The Sentry organization, used to send a direct link from PostHog to Sentry
-    project_id = None  # The Sentry project id, used to send a direct link from PostHog to Sentry
+    project_id = (
+        None  # The Sentry project id, used to send a direct link from PostHog to Sentry
+    )
     prefix = "https://sentry.io/organizations/"  # URL of a hosted sentry instance (default: https://sentry.io/organizations/)
 
     @staticmethod
@@ -31,7 +33,9 @@ class PostHogIntegration(Integration):
 
                 if event.get("tags", {}).get(POSTHOG_ID_TAG):
                     posthog_distinct_id = event["tags"][POSTHOG_ID_TAG]
-                    event["tags"]["PostHog URL"] = f"{posthog.host or DEFAULT_HOST}/person/{posthog_distinct_id}"
+                    event["tags"]["PostHog URL"] = (
+                        f"{posthog.host or DEFAULT_HOST}/person/{posthog_distinct_id}"
+                    )
 
                     properties = {
                         "$sentry_event_id": event["event_id"],
@@ -40,7 +44,8 @@ class PostHogIntegration(Integration):
 
                     if PostHogIntegration.organization:
                         project_id = PostHogIntegration.project_id or (
-                            not not Hub.current.client.dsn and Dsn(Hub.current.client.dsn).project_id
+                            not not Hub.current.client.dsn
+                            and Dsn(Hub.current.client.dsn).project_id
                         )
                         if project_id:
                             properties["$sentry_url"] = (

posthog/test/ai/anthropic/test_anthropic.py

@@ -16,7 +16,9 @@ except ImportError:
 ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY")
 
 # Skip all tests if Anthropic is not available
-pytestmark = pytest.mark.skipif(not ANTHROPIC_AVAILABLE, reason="Anthropic package is not available")
+pytestmark = pytest.mark.skipif(
+    not ANTHROPIC_AVAILABLE, reason="Anthropic package is not available"
+)
 
 
 @pytest.fixture
@@ -87,7 +89,9 @@ def mock_anthropic_response_with_cached_tokens():
 
 
 def test_basic_completion(mock_client, mock_anthropic_response):
-    with patch("anthropic.resources.Messages.create", return_value=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",
@@ -107,7 +111,9 @@ def test_basic_completion(mock_client, mock_anthropic_response):
         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_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
@@ -116,7 +122,9 @@ def test_basic_completion(mock_client, mock_anthropic_response):
 
 
 def test_streaming(mock_client, mock_anthropic_stream):
-    with patch("anthropic.resources.Messages.create", return_value=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",
@@ -153,7 +161,9 @@ def test_streaming(mock_client, mock_anthropic_stream):
 
 
 def test_streaming_with_stream_endpoint(mock_client, mock_anthropic_stream):
-    with patch("anthropic.resources.Messages.create", return_value=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",
@@ -189,7 +199,9 @@ def test_streaming_with_stream_endpoint(mock_client, mock_anthropic_stream):
 
 
 def test_groups(mock_client, mock_anthropic_response):
-    with patch("anthropic.resources.Messages.create", return_value=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",
@@ -206,7 +218,9 @@ def test_groups(mock_client, mock_anthropic_response):
 
 
 def test_privacy_mode_local(mock_client, mock_anthropic_response):
-    with patch("anthropic.resources.Messages.create", return_value=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",
@@ -225,7 +239,9 @@ def test_privacy_mode_local(mock_client, mock_anthropic_response):
 
 
 def test_privacy_mode_global(mock_client, mock_anthropic_response):
-    with patch("anthropic.resources.Messages.create", return_value=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(
@@ -299,7 +315,9 @@ async def test_basic_async_integration(mock_client):
     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_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
@@ -309,7 +327,9 @@ async def test_basic_async_integration(mock_client):
 
 
 def test_streaming_system_prompt(mock_client, mock_anthropic_stream):
-    with patch("anthropic.resources.Messages.create", return_value=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",
@@ -328,7 +348,10 @@ def test_streaming_system_prompt(mock_client, mock_anthropic_stream):
         call_args = mock_client.capture.call_args[1]
         props = call_args["properties"]
 
-        assert props["$ai_input"] == [{"role": "system", "content": "Foo"}, {"role": "user", "content": "Bar"}]
+        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")
@@ -359,10 +382,15 @@ async def test_async_streaming_system_prompt(mock_client, mock_anthropic_stream)
 
 
 def test_error(mock_client, mock_anthropic_response):
-    with patch("anthropic.resources.Messages.create", side_effect=Exception("Test error")):
+    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"}])
+            client.messages.create(
+                model="claude-3-opus-20240229",
+                messages=[{"role": "user", "content": "Hello"}],
+            )
 
         assert mock_client.capture.call_count == 1
 
@@ -373,7 +401,10 @@ def test_error(mock_client, mock_anthropic_response):
 
 
 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):
+    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",
@@ -393,7 +424,9 @@ def test_cached_tokens(mock_client, mock_anthropic_response_with_cached_tokens):
         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_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

posthog/test/ai/gemini/test_gemini.py

@@ -0,0 +1,320 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+try:
+    from google import genai as google_genai
+
+    from posthog.ai.gemini import Client
+
+    GEMINI_AVAILABLE = True
+except ImportError:
+    GEMINI_AVAILABLE = False
+
+pytestmark = pytest.mark.skipif(
+    not GEMINI_AVAILABLE, reason="Google Gemini package is not available"
+)
+
+
+@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_gemini_response():
+    mock_response = MagicMock()
+    mock_response.text = "Test response from Gemini"
+
+    mock_usage = MagicMock()
+    mock_usage.prompt_token_count = 20
+    mock_usage.candidates_token_count = 10
+    mock_response.usage_metadata = mock_usage
+
+    mock_candidate = MagicMock()
+    mock_candidate.text = "Test response from Gemini"
+    mock_content = MagicMock()
+    mock_part = MagicMock()
+    mock_part.text = "Test response from Gemini"
+    mock_content.parts = [mock_part]
+    mock_candidate.content = mock_content
+    mock_response.candidates = [mock_candidate]
+
+    return mock_response
+
+
+@pytest.fixture
+def mock_google_genai_client():
+    """Mock for the new google-genai Client"""
+    with patch.object(google_genai, "Client") as mock_client_class:
+        mock_client_instance = MagicMock()
+        mock_models = MagicMock()
+        mock_client_instance.models = mock_models
+        mock_client_class.return_value = mock_client_instance
+        yield mock_client_instance
+
+
+def test_new_client_basic_generation(
+    mock_client, mock_google_genai_client, mock_gemini_response
+):
+    """Test the new Client/Models API structure"""
+    mock_google_genai_client.models.generate_content.return_value = mock_gemini_response
+
+    client = Client(api_key="test-key", posthog_client=mock_client)
+
+    response = client.models.generate_content(
+        model="gemini-2.0-flash",
+        contents=["Tell me a fun fact about hedgehogs"],
+        posthog_distinct_id="test-id",
+        posthog_properties={"foo": "bar"},
+    )
+
+    assert response == mock_gemini_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"] == "gemini"
+    assert props["$ai_model"] == "gemini-2.0-flash"
+    assert props["$ai_input_tokens"] == 20
+    assert props["$ai_output_tokens"] == 10
+    assert props["foo"] == "bar"
+    assert "$ai_trace_id" in props
+    assert props["$ai_latency"] > 0
+
+
+def test_new_client_streaming_with_generate_content_stream(
+    mock_client, mock_google_genai_client
+):
+    """Test the new generate_content_stream method"""
+
+    def mock_streaming_response():
+        mock_chunk1 = MagicMock()
+        mock_chunk1.text = "Hello "
+        mock_usage1 = MagicMock()
+        mock_usage1.prompt_token_count = 10
+        mock_usage1.candidates_token_count = 5
+        mock_chunk1.usage_metadata = mock_usage1
+
+        mock_chunk2 = MagicMock()
+        mock_chunk2.text = "world!"
+        mock_usage2 = MagicMock()
+        mock_usage2.prompt_token_count = 10
+        mock_usage2.candidates_token_count = 10
+        mock_chunk2.usage_metadata = mock_usage2
+
+        yield mock_chunk1
+        yield mock_chunk2
+
+    # Mock the generate_content_stream method
+    mock_google_genai_client.models.generate_content_stream.return_value = (
+        mock_streaming_response()
+    )
+
+    client = Client(api_key="test-key", posthog_client=mock_client)
+
+    response = client.models.generate_content_stream(
+        model="gemini-2.0-flash",
+        contents=["Write a short story"],
+        posthog_distinct_id="test-id",
+        posthog_properties={"feature": "streaming"},
+    )
+
+    chunks = list(response)
+    assert len(chunks) == 2
+    assert chunks[0].text == "Hello "
+    assert chunks[1].text == "world!"
+
+    # Check that the streaming event was captured
+    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"] == "gemini"
+    assert props["$ai_model"] == "gemini-2.0-flash"
+    assert props["$ai_input_tokens"] == 10
+    assert props["$ai_output_tokens"] == 10
+    assert props["feature"] == "streaming"
+    assert isinstance(props["$ai_latency"], float)
+
+
+def test_new_client_groups(mock_client, mock_google_genai_client, mock_gemini_response):
+    """Test groups functionality with new Client API"""
+    mock_google_genai_client.models.generate_content.return_value = mock_gemini_response
+
+    client = Client(api_key="test-key", posthog_client=mock_client)
+
+    client.models.generate_content(
+        model="gemini-2.0-flash",
+        contents=["Hello"],
+        posthog_distinct_id="test-id",
+        posthog_groups={"company": "company_123"},
+    )
+
+    call_args = mock_client.capture.call_args[1]
+    assert call_args["groups"] == {"company": "company_123"}
+
+
+def test_new_client_privacy_mode_local(
+    mock_client, mock_google_genai_client, mock_gemini_response
+):
+    """Test local privacy mode with new Client API"""
+    mock_google_genai_client.models.generate_content.return_value = mock_gemini_response
+
+    client = Client(api_key="test-key", posthog_client=mock_client)
+
+    client.models.generate_content(
+        model="gemini-2.0-flash",
+        contents=["Hello"],
+        posthog_distinct_id="test-id",
+        posthog_privacy_mode=True,
+    )
+
+    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_new_client_privacy_mode_global(
+    mock_client, mock_google_genai_client, mock_gemini_response
+):
+    """Test global privacy mode with new Client API"""
+    mock_client.privacy_mode = True
+
+    mock_google_genai_client.models.generate_content.return_value = mock_gemini_response
+
+    client = Client(api_key="test-key", posthog_client=mock_client)
+
+    client.models.generate_content(
+        model="gemini-2.0-flash",
+        contents=["Hello"],
+        posthog_distinct_id="test-id",
+    )
+
+    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_new_client_different_input_formats(
+    mock_client, mock_google_genai_client, mock_gemini_response
+):
+    """Test different input formats with new Client API"""
+    mock_google_genai_client.models.generate_content.return_value = mock_gemini_response
+
+    client = Client(api_key="test-key", posthog_client=mock_client)
+
+    # Test string input
+    client.models.generate_content(
+        model="gemini-2.0-flash", contents="Hello", posthog_distinct_id="test-id"
+    )
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+    assert props["$ai_input"] == [{"role": "user", "content": "Hello"}]
+
+    # Test list input
+    mock_client.capture.reset_mock()
+    mock_part = MagicMock()
+    mock_part.text = "List item"
+    client.models.generate_content(
+        model="gemini-2.0-flash", contents=[mock_part], posthog_distinct_id="test-id"
+    )
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+    assert props["$ai_input"] == [{"role": "user", "content": "List item"}]
+
+
+def test_new_client_model_parameters(
+    mock_client, mock_google_genai_client, mock_gemini_response
+):
+    """Test model parameters with new Client API"""
+    mock_google_genai_client.models.generate_content.return_value = mock_gemini_response
+
+    client = Client(api_key="test-key", posthog_client=mock_client)
+
+    client.models.generate_content(
+        model="gemini-2.0-flash",
+        contents=["Hello"],
+        posthog_distinct_id="test-id",
+        temperature=0.7,
+        max_tokens=100,
+    )
+
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+    assert props["$ai_model_parameters"]["temperature"] == 0.7
+    assert props["$ai_model_parameters"]["max_tokens"] == 100
+
+
+def test_new_client_default_settings(
+    mock_client, mock_google_genai_client, mock_gemini_response
+):
+    """Test client with default PostHog settings"""
+    mock_google_genai_client.models.generate_content.return_value = mock_gemini_response
+
+    client = Client(
+        api_key="test-key",
+        posthog_client=mock_client,
+        posthog_distinct_id="default_user",
+        posthog_properties={"team": "ai"},
+        posthog_privacy_mode=False,
+        posthog_groups={"company": "acme_corp"},
+    )
+
+    # Call without overriding defaults
+    client.models.generate_content(model="gemini-2.0-flash", contents=["Hello"])
+
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+
+    assert call_args["distinct_id"] == "default_user"
+    assert call_args["groups"] == {"company": "acme_corp"}
+    assert props["team"] == "ai"
+
+
+def test_new_client_override_defaults(
+    mock_client, mock_google_genai_client, mock_gemini_response
+):
+    """Test overriding client defaults per call"""
+    mock_google_genai_client.models.generate_content.return_value = mock_gemini_response
+
+    client = Client(
+        api_key="test-key",
+        posthog_client=mock_client,
+        posthog_distinct_id="default_user",
+        posthog_properties={"team": "ai"},
+        posthog_privacy_mode=False,
+        posthog_groups={"company": "acme_corp"},
+    )
+
+    # Override defaults in call
+    client.models.generate_content(
+        model="gemini-2.0-flash",
+        contents=["Hello"],
+        posthog_distinct_id="specific_user",
+        posthog_properties={"feature": "chat", "urgent": True},
+        posthog_privacy_mode=True,
+        posthog_groups={"organization": "special_org"},
+    )
+
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+
+    # Check overrides
+    assert call_args["distinct_id"] == "specific_user"
+    assert call_args["groups"] == {"organization": "special_org"}
+    assert props["$ai_input"] is None  # privacy mode was overridden
+
+    # Check merged properties (defaults + call-specific)
+    assert props["team"] == "ai"  # from defaults
+    assert props["feature"] == "chat"  # from call
+    assert props["urgent"] is True  # from call

posthog/test/ai/langchain/test_callbacks.py

@@ -43,7 +43,9 @@ except ImportError:
 
 
 # Skip all tests if LangChain is not available
-pytestmark = pytest.mark.skipif(not LANGCHAIN_AVAILABLE, reason="LangChain package is not available")
+pytestmark = pytest.mark.skipif(
+    not LANGCHAIN_AVAILABLE, reason="LangChain package is not available"
+)
 
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY")
@@ -252,7 +254,9 @@ async def test_async_basic_chat_chain(mock_client, stream):
     callbacks = [CallbackHandler(mock_client)]
     chain = prompt | model
     if stream:
-        result = [m async for m in chain.astream({}, config={"callbacks": callbacks})][0]
+        result = [m async for m in chain.astream({}, config={"callbacks": callbacks})][
+            0
+        ]
     else:
         result = await chain.ainvoke({}, config={"callbacks": callbacks})
     assert result.content == "The Los Angeles Dodgers won the World Series in 2020."
@@ -316,10 +320,17 @@ def test_basic_llm_chain(mock_client, Model, stream):
 
     if stream:
         result = "".join(
-            [m for m in model.stream("Who won the world series in 2020?", config={"callbacks": callbacks})]
+            [
+                m
+                for m in model.stream(
+                    "Who won the world series in 2020?", config={"callbacks": callbacks}
+                )
+            ]
         )
     else:
-        result = model.invoke("Who won the world series in 2020?", config={"callbacks": callbacks})
+        result = model.invoke(
+            "Who won the world series in 2020?", config={"callbacks": callbacks}
+        )
     assert result == "The Los Angeles Dodgers won the World Series in 2020."
 
     assert mock_client.capture.call_count == 1
@@ -331,7 +342,9 @@ def test_basic_llm_chain(mock_client, Model, stream):
     assert "$ai_model" in props
     assert "$ai_provider" in props
     assert props["$ai_input"] == ["Who won the world series in 2020?"]
-    assert props["$ai_output_choices"] == ["The Los Angeles Dodgers won the World Series in 2020."]
+    assert props["$ai_output_choices"] == [
+        "The Los Angeles Dodgers won the World Series in 2020."
+    ]
     assert props["$ai_http_status"] == 200
     assert props["$ai_trace_id"] is not None
     assert isinstance(props["$ai_latency"], float)
@@ -352,10 +365,17 @@ async def test_async_basic_llm_chain(mock_client, Model, stream):
 
     if stream:
         result = "".join(
-            [m async for m in model.astream("Who won the world series in 2020?", config={"callbacks": callbacks})]
+            [
+                m
+                async for m in model.astream(
+                    "Who won the world series in 2020?", config={"callbacks": callbacks}
+                )
+            ]
         )
     else:
-        result = await model.ainvoke("Who won the world series in 2020?", config={"callbacks": callbacks})
+        result = await model.ainvoke(
+            "Who won the world series in 2020?", config={"callbacks": callbacks}
+        )
     assert result == "The Los Angeles Dodgers won the World Series in 2020."
 
     assert mock_client.capture.call_count == 1
@@ -367,7 +387,9 @@ async def test_async_basic_llm_chain(mock_client, Model, stream):
     assert "$ai_model" in props
     assert "$ai_provider" in props
     assert props["$ai_input"] == ["Who won the world series in 2020?"]
-    assert props["$ai_output_choices"] == ["The Los Angeles Dodgers won the World Series in 2020."]
+    assert props["$ai_output_choices"] == [
+        "The Los Angeles Dodgers won the World Series in 2020."
+    ]
     assert props["$ai_http_status"] == 200
     assert props["$ai_trace_id"] is not None
     assert isinstance(props["$ai_latency"], float)
@@ -418,8 +440,12 @@ def test_trace_id_and_inputs_for_multiple_chains(mock_client):
     assert "distinct_id" in first_generation_args
     assert "$ai_model" in first_generation_props
     assert "$ai_provider" in first_generation_props
-    assert first_generation_props["$ai_input"] == [{"role": "user", "content": "Foo bar"}]
-    assert first_generation_props["$ai_output_choices"] == [{"role": "assistant", "content": "Bar"}]
+    assert first_generation_props["$ai_input"] == [
+        {"role": "user", "content": "Foo bar"}
+    ]
+    assert first_generation_props["$ai_output_choices"] == [
+        {"role": "assistant", "content": "Bar"}
+    ]
     assert first_generation_props["$ai_http_status"] == 200
     assert isinstance(first_generation_props["$ai_latency"], float)
     assert "$ai_span_id" in first_generation_props
@@ -439,8 +465,12 @@ def test_trace_id_and_inputs_for_multiple_chains(mock_client):
     assert "distinct_id" in second_generation_args
     assert "$ai_model" in second_generation_props
     assert "$ai_provider" in second_generation_props
-    assert second_generation_props["$ai_input"] == [{"role": "assistant", "content": "Bar"}]
-    assert second_generation_props["$ai_output_choices"] == [{"role": "assistant", "content": "Bar"}]
+    assert second_generation_props["$ai_input"] == [
+        {"role": "assistant", "content": "Bar"}
+    ]
+    assert second_generation_props["$ai_output_choices"] == [
+        {"role": "assistant", "content": "Bar"}
+    ]
     assert second_generation_props["$ai_http_status"] == 200
     assert second_generation_props["$ai_trace_id"] is not None
     assert isinstance(second_generation_props["$ai_latency"], float)
@@ -475,7 +505,9 @@ def test_personless_mode(mock_client):
     assert trace_args["properties"]["$process_person_profile"] is False
 
     id = uuid.uuid4()
-    chain.invoke({}, config={"callbacks": [CallbackHandler(mock_client, distinct_id=id)]})
+    chain.invoke(
+        {}, config={"callbacks": [CallbackHandler(mock_client, distinct_id=id)]}
+    )
     assert mock_client.capture.call_count == 6
     span_args = mock_client.capture.call_args_list[3][1]
     generation_args = mock_client.capture.call_args_list[4][1]
@@ -515,7 +547,9 @@ def test_personless_mode_exception(mock_client):
 
     id = uuid.uuid4()
     with pytest.raises(Exception):
-        chain.invoke({}, config={"callbacks": [CallbackHandler(mock_client, distinct_id=id)]})
+        chain.invoke(
+            {}, config={"callbacks": [CallbackHandler(mock_client, distinct_id=id)]}
+        )
     assert mock_client.capture.call_count == 6
     span_args = mock_client.capture.call_args_list[3][1]
     generation_args = mock_client.capture.call_args_list[4][1]
@@ -574,7 +608,9 @@ def test_metadata(mock_client):
     assert generation_call_props["$ai_trace_id"] == "test-trace-id"
     assert generation_call_props["foo"] == "bar"
     assert generation_call_props["$ai_input"] == [{"role": "user", "content": "Foo"}]
-    assert generation_call_props["$ai_output_choices"] == [{"role": "assistant", "content": "Bar"}]
+    assert generation_call_props["$ai_output_choices"] == [
+        {"role": "assistant", "content": "Bar"}
+    ]
     assert generation_call_props["$ai_http_status"] == 200
     assert isinstance(generation_call_props["$ai_latency"], float)
 
@@ -655,7 +691,10 @@ def test_graph_state(mock_client):
 
     # 1. Span, finish initialization
     second_state = {
-        "messages": [HumanMessage(content="What's a bar?"), AIMessage(content="Let's explore bar.")],
+        "messages": [
+            HumanMessage(content="What's a bar?"),
+            AIMessage(content="Let's explore bar."),
+        ],
         "xyz": "abc",
     }
 
@@ -669,19 +708,25 @@ def test_graph_state(mock_client):
 
     # 2. Span - the ChatPromptTemplate within fake_llm's FakeMessagesListChatModel
     assert calls[1]["event"] == "$ai_span"
-    assert calls[1]["properties"]["$ai_parent_id"] == calls[3]["properties"]["$ai_span_id"]
+    assert (
+        calls[1]["properties"]["$ai_parent_id"] == calls[3]["properties"]["$ai_span_id"]
+    )
     assert "$ai_span_id" in calls[1]["properties"]
     assert calls[1]["properties"]["$ai_span_name"] == "ChatPromptTemplate"
 
     # 3. Generation - the FakeMessagesListChatModel within fake_llm's RunnableSequence
     assert calls[2]["event"] == "$ai_generation"
-    assert calls[2]["properties"]["$ai_parent_id"] == calls[3]["properties"]["$ai_span_id"]
+    assert (
+        calls[2]["properties"]["$ai_parent_id"] == calls[3]["properties"]["$ai_span_id"]
+    )
     assert "$ai_span_id" in calls[2]["properties"]
     assert calls[2]["properties"]["$ai_span_name"] == "FakeMessagesListChatModel"
 
     # 4. Span - RunnableSequence within fake_llm
     assert calls[3]["event"] == "$ai_span"
-    assert calls[3]["properties"]["$ai_parent_id"] == calls[4]["properties"]["$ai_span_id"]
+    assert (
+        calls[3]["properties"]["$ai_parent_id"] == calls[4]["properties"]["$ai_span_id"]
+    )
     assert "$ai_span_id" in calls[3]["properties"]
     assert calls[3]["properties"]["$ai_span_name"] == "RunnableSequence"
 
@@ -705,9 +750,14 @@ def test_graph_state(mock_client):
     assert isinstance(trace_props["$ai_output_state"]["messages"][0], HumanMessage)
     assert trace_props["$ai_output_state"]["messages"][0].content == "What's a bar?"
     assert isinstance(trace_props["$ai_output_state"]["messages"][1], AIMessage)
-    assert trace_props["$ai_output_state"]["messages"][1].content == "Let's explore bar."
+    assert (
+        trace_props["$ai_output_state"]["messages"][1].content == "Let's explore bar."
+    )
     assert isinstance(trace_props["$ai_output_state"]["messages"][2], AIMessage)
-    assert trace_props["$ai_output_state"]["messages"][2].content == "It's a type of greeble."
+    assert (
+        trace_props["$ai_output_state"]["messages"][2].content
+        == "It's a type of greeble."
+    )
     assert trace_args["properties"]["$ai_output_state"]["xyz"] == "abc"
 
 
@@ -735,7 +785,9 @@ def test_callbacks_logic(mock_client):
         assert len(callbacks._parent_tree.items()) == 1
         return [m]
 
-    (chain | RunnableLambda(assert_intermediary_run) | model).invoke({}, config={"callbacks": [callbacks]})
+    (chain | RunnableLambda(assert_intermediary_run) | model).invoke(
+        {}, config={"callbacks": [callbacks]}
+    )
     assert callbacks._runs == {}
     assert callbacks._parent_tree == {}
 
@@ -828,10 +880,16 @@ def test_openai_chain(mock_client):
         {"role": "system", "content": 'You must always answer with "Bar".'},
         {"role": "user", "content": "Foo"},
     ]
-    assert gen_props["$ai_output_choices"] == [{"role": "assistant", "content": "Bar", "refusal": None}]
+    assert gen_props["$ai_output_choices"] == [
+        {"role": "assistant", "content": "Bar", "refusal": None}
+    ]
     assert gen_props["$ai_http_status"] == 200
     assert isinstance(gen_props["$ai_latency"], float)
-    assert min(approximate_latency - 1, 0) <= math.floor(gen_props["$ai_latency"]) <= approximate_latency
+    assert (
+        min(approximate_latency - 1, 0)
+        <= math.floor(gen_props["$ai_latency"])
+        <= approximate_latency
+    )
     assert gen_props["$ai_input_tokens"] == 20
     assert gen_props["$ai_output_tokens"] == 1
 
@@ -1105,7 +1163,11 @@ def test_anthropic_chain(mock_client):
     assert gen_props["$ai_output_choices"] == [{"role": "assistant", "content": "Bar"}]
     assert gen_props["$ai_http_status"] == 200
     assert isinstance(gen_props["$ai_latency"], float)
-    assert min(approximate_latency - 1, 0) <= math.floor(gen_props["$ai_latency"]) <= approximate_latency
+    assert (
+        min(approximate_latency - 1, 0)
+        <= math.floor(gen_props["$ai_latency"])
+        <= approximate_latency
+    )
     assert gen_props["$ai_input_tokens"] == 17
     assert gen_props["$ai_output_tokens"] == 1
 
@@ -1252,7 +1314,10 @@ def test_tool_calls(mock_client):
             },
         }
     ]
-    assert "additional_kwargs" not in generation_call["properties"]["$ai_output_choices"][0]
+    assert (
+        "additional_kwargs"
+        not in generation_call["properties"]["$ai_output_choices"][0]
+    )
 
 
 async def test_async_traces(mock_client):
@@ -1274,7 +1339,9 @@ async def test_async_traces(mock_client):
     approximate_latency = math.floor(time.time() - start_time)
     assert mock_client.capture.call_count == 4
 
-    first_call, second_call, third_call, fourth_call = mock_client.capture.call_args_list
+    first_call, second_call, third_call, fourth_call = (
+        mock_client.capture.call_args_list
+    )
     assert first_call[1]["event"] == "$ai_span"
     assert second_call[1]["event"] == "$ai_generation"
     assert third_call[1]["event"] == "$ai_trace"
@@ -1282,7 +1349,9 @@ async def test_async_traces(mock_client):
     assert fourth_call[1]["event"] == "$ai_trace"
     assert fourth_call[1]["properties"]["$ai_span_name"] == "sleep"
     assert (
-        min(approximate_latency - 1, 0) <= math.floor(third_call[1]["properties"]["$ai_latency"]) <= approximate_latency
+        min(approximate_latency - 1, 0)
+        <= math.floor(third_call[1]["properties"]["$ai_latency"])
+        <= approximate_latency
     )
 
 
@@ -1304,7 +1373,9 @@ def test_langgraph_agent(mock_client):
     model = ChatOpenAI(api_key=OPENAI_API_KEY, model="gpt-4o-mini", temperature=0)
     graph = create_react_agent(model, tools=tools)
     inputs = {"messages": [("user", "what is the weather in sf")]}
-    cb = CallbackHandler(mock_client, trace_id="test-trace-id", distinct_id="test-distinct-id")
+    cb = CallbackHandler(
+        mock_client, trace_id="test-trace-id", distinct_id="test-distinct-id"
+    )
     graph.invoke(inputs, config={"callbacks": [cb]})
     calls = [call[1] for call in mock_client.capture.call_args_list]
     assert len(calls) == 21
@@ -1324,7 +1395,9 @@ def test_span_set_parent_ids(mock_client, trace_id):
         ]
     )
     model = FakeMessagesListChatModel(
-        responses=[AIMessage(content="The Los Angeles Dodgers won the World Series in 2020.")]
+        responses=[
+            AIMessage(content="The Los Angeles Dodgers won the World Series in 2020.")
+        ]
     )
     callbacks = [CallbackHandler(mock_client, trace_id=trace_id)]
     chain = prompt | model
@@ -1333,10 +1406,16 @@ def test_span_set_parent_ids(mock_client, trace_id):
     assert mock_client.capture.call_count == 3
 
     span_props = mock_client.capture.call_args_list[0][1]
-    assert span_props["properties"]["$ai_trace_id"] == span_props["properties"]["$ai_parent_id"]
+    assert (
+        span_props["properties"]["$ai_trace_id"]
+        == span_props["properties"]["$ai_parent_id"]
+    )
 
     generation_props = mock_client.capture.call_args_list[1][1]
-    assert generation_props["properties"]["$ai_trace_id"] == generation_props["properties"]["$ai_parent_id"]
+    assert (
+        generation_props["properties"]["$ai_trace_id"]
+        == generation_props["properties"]["$ai_parent_id"]
+    )
 
 
 @pytest.mark.parametrize("trace_id", ["test-trace-id", None])
@@ -1356,7 +1435,9 @@ def test_span_set_parent_ids_for_third_level_run(mock_client, trace_id):
 
     assert mock_client.capture.call_count == 3
 
-    span2, span1, trace = [call[1]["properties"] for call in mock_client.capture.call_args_list]
+    span2, span1, trace = [
+        call[1]["properties"] for call in mock_client.capture.call_args_list
+    ]
     assert span2["$ai_parent_id"] == span1["$ai_span_id"]
     assert span1["$ai_parent_id"] == trace["$ai_trace_id"]
 
@@ -1373,7 +1454,10 @@ def test_captures_error_with_details_in_span(mock_client):
         pass
 
     assert mock_client.capture.call_count == 2
-    assert mock_client.capture.call_args_list[1][1]["properties"]["$ai_error"] == "ValueError: test"
+    assert (
+        mock_client.capture.call_args_list[1][1]["properties"]["$ai_error"]
+        == "ValueError: test"
+    )
     assert mock_client.capture.call_args_list[1][1]["properties"]["$ai_is_error"]
 
 
@@ -1389,5 +1473,8 @@ def test_captures_error_without_details_in_span(mock_client):
         pass
 
     assert mock_client.capture.call_count == 2
-    assert mock_client.capture.call_args_list[1][1]["properties"]["$ai_error"] == "ValueError"
+    assert (
+        mock_client.capture.call_args_list[1][1]["properties"]["$ai_error"]
+        == "ValueError"
+    )
     assert mock_client.capture.call_args_list[1][1]["properties"]["$ai_is_error"]

posthog/test/ai/openai/test_openai.py

@@ -9,12 +9,24 @@ try:
     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.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 openai.types.responses import (
+        Response,
+        ResponseOutputMessage,
+        ResponseOutputText,
+        ResponseUsage,
+    )
 
     from posthog.ai.openai import OpenAI
 
@@ -23,7 +35,9 @@ except ImportError:
     OPENAI_AVAILABLE = False
 
 # Skip all tests if OpenAI is not available
-pytestmark = pytest.mark.skipif(not OPENAI_AVAILABLE, reason="OpenAI package is not available")
+pytestmark = pytest.mark.skipif(
+    not OPENAI_AVAILABLE, reason="OpenAI package is not available"
+)
 
 
 @pytest.fixture
@@ -182,7 +196,10 @@ def mock_openai_response_with_tool_calls():
 
 
 def test_basic_completion(mock_client, mock_openai_response):
-    with patch("openai.resources.chat.completions.Completions.create", return_value=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",
@@ -202,7 +219,9 @@ def test_basic_completion(mock_client, mock_openai_response):
         assert props["$ai_provider"] == "openai"
         assert props["$ai_model"] == "gpt-4"
         assert props["$ai_input"] == [{"role": "user", "content": "Hello"}]
-        assert props["$ai_output_choices"] == [{"role": "assistant", "content": "Test response"}]
+        assert props["$ai_output_choices"] == [
+            {"role": "assistant", "content": "Test response"}
+        ]
         assert props["$ai_input_tokens"] == 20
         assert props["$ai_output_tokens"] == 10
         assert props["$ai_http_status"] == 200
@@ -211,7 +230,10 @@ def test_basic_completion(mock_client, mock_openai_response):
 
 
 def test_embeddings(mock_client, mock_embedding_response):
-    with patch("openai.resources.embeddings.Embeddings.create", return_value=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",
@@ -238,7 +260,10 @@ def test_embeddings(mock_client, mock_embedding_response):
 
 
 def test_groups(mock_client, mock_openai_response):
-    with patch("openai.resources.chat.completions.Completions.create", return_value=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",
@@ -256,7 +281,10 @@ def test_groups(mock_client, mock_openai_response):
 
 
 def test_privacy_mode_local(mock_client, mock_openai_response):
-    with patch("openai.resources.chat.completions.Completions.create", return_value=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",
@@ -275,7 +303,10 @@ def test_privacy_mode_local(mock_client, mock_openai_response):
 
 
 def test_privacy_mode_global(mock_client, mock_openai_response):
-    with patch("openai.resources.chat.completions.Completions.create", return_value=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(
@@ -295,10 +326,15 @@ def test_privacy_mode_global(mock_client, mock_openai_response):
 
 
 def test_error(mock_client, mock_openai_response):
-    with patch("openai.resources.chat.completions.Completions.create", side_effect=Exception("Test error")):
+    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"}])
+            client.chat.completions.create(
+                model="gpt-4", messages=[{"role": "user", "content": "Hello"}]
+            )
 
         assert mock_client.capture.call_count == 1
 
@@ -310,7 +346,8 @@ def test_error(mock_client, mock_openai_response):
 
 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
+        "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(
@@ -331,7 +368,9 @@ def test_cached_tokens(mock_client, mock_openai_response_with_cached_tokens):
         assert props["$ai_provider"] == "openai"
         assert props["$ai_model"] == "gpt-4"
         assert props["$ai_input"] == [{"role": "user", "content": "Hello"}]
-        assert props["$ai_output_choices"] == [{"role": "assistant", "content": "Test response"}]
+        assert props["$ai_output_choices"] == [
+            {"role": "assistant", "content": "Test response"}
+        ]
         assert props["$ai_input_tokens"] == 20
         assert props["$ai_output_tokens"] == 10
         assert props["$ai_cache_read_input_tokens"] == 15
@@ -342,16 +381,23 @@ def test_cached_tokens(mock_client, mock_openai_response_with_cached_tokens):
 
 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
+        "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?"}],
+            messages=[
+                {"role": "user", "content": "What's the weather in San Francisco?"}
+            ],
             tools=[
                 {
                     "type": "function",
-                    "function": {"name": "get_weather", "description": "Get weather", "parameters": {}},
+                    "function": {
+                        "name": "get_weather",
+                        "description": "Get weather",
+                        "parameters": {},
+                    },
                 }
             ],
             posthog_distinct_id="test-id",
@@ -367,8 +413,12 @@ def test_tool_calls(mock_client, mock_openai_response_with_tool_calls):
         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."}]
+        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
@@ -501,11 +551,17 @@ def test_streaming_with_tool_calls(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?"}],
+            messages=[
+                {"role": "user", "content": "What's the weather in San Francisco?"}
+            ],
             tools=[
                 {
                     "type": "function",
-                    "function": {"name": "get_weather", "description": "Get weather", "parameters": {}},
+                    "function": {
+                        "name": "get_weather",
+                        "description": "Get weather",
+                        "parameters": {},
+                    },
                 }
             ],
             stream=True,
@@ -547,7 +603,10 @@ def test_streaming_with_tool_calls(mock_client):
         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."
+        assert (
+            props["$ai_output_choices"][0]["content"]
+            == "The weather in San Francisco is 15°C."
+        )
 
         # Check token usage
         assert props["$ai_input_tokens"] == 20
@@ -556,7 +615,10 @@ def test_streaming_with_tool_calls(mock_client):
 
 # 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):
+    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",
@@ -575,7 +637,9 @@ def test_responses_api(mock_client, mock_openai_response_with_responses_api):
         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_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

posthog/test/exception_integrations/test_django.py

@@ -1,8 +1,6 @@
 from posthog.exception_integrations.django import DjangoRequestExtractor
 
-DEFAULT_USER_AGENT = (
-    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3"
-)
+DEFAULT_USER_AGENT = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3"
 
 
 def mock_request_factory(override_headers):
@@ -31,7 +29,9 @@ def test_request_extractor_with_no_trace():
 
 
 def test_request_extractor_with_trace():
-    request = mock_request_factory({"traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"})
+    request = mock_request_factory(
+        {"traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"}
+    )
     extractor = DjangoRequestExtractor(request)
     assert extractor.extract_person_data() == {
         "ip": "193.4.5.12",
@@ -58,7 +58,9 @@ def test_request_extractor_with_tracestate():
 
 
 def test_request_extractor_with_complicated_tracestate():
-    request = mock_request_factory({"tracestate": "posthog-distinct-id=alohaMountainsXUYZ,rojo=00f067aa0ba902b7"})
+    request = mock_request_factory(
+        {"tracestate": "posthog-distinct-id=alohaMountainsXUYZ,rojo=00f067aa0ba902b7"}
+    )
     extractor = DjangoRequestExtractor(request)
     assert extractor.extract_person_data() == {
         "ip": "193.4.5.12",

posthog/test/test_client.py

@@ -164,7 +164,9 @@ 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")
+            client = Client(
+                FAKE_TEST_API_KEY, on_error=self.set_fail, host="https://aloha.com"
+            )
             exception = Exception("test exception")
             client.capture_exception(exception, "distinct_id")
 
@@ -189,9 +191,15 @@ class TestClient(unittest.TestCase):
                 },
             )
 
-    def test_basic_capture_exception_with_correct_host_generation_for_server_hosts(self):
+    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")
+            client = Client(
+                FAKE_TEST_API_KEY,
+                on_error=self.set_fail,
+                host="https://app.posthog.com",
+            )
             exception = Exception("test exception")
             client.capture_exception(exception, "distinct_id")
 
@@ -230,27 +238,45 @@ class TestClient(unittest.TestCase):
             self.assertEqual(capture_call[1], "$exception")
             self.assertEqual(capture_call[2]["$exception_type"], "Exception")
             self.assertEqual(capture_call[2]["$exception_message"], "test exception")
-            self.assertEqual(capture_call[2]["$exception_list"][0]["mechanism"]["type"], "generic")
-            self.assertEqual(capture_call[2]["$exception_list"][0]["mechanism"]["handled"], True)
+            self.assertEqual(
+                capture_call[2]["$exception_list"][0]["mechanism"]["type"], "generic"
+            )
+            self.assertEqual(
+                capture_call[2]["$exception_list"][0]["mechanism"]["handled"], True
+            )
             self.assertEqual(capture_call[2]["$exception_list"][0]["module"], None)
             self.assertEqual(capture_call[2]["$exception_list"][0]["type"], "Exception")
-            self.assertEqual(capture_call[2]["$exception_list"][0]["value"], "test exception")
+            self.assertEqual(
+                capture_call[2]["$exception_list"][0]["value"], "test exception"
+            )
             self.assertEqual(
                 capture_call[2]["$exception_list"][0]["stacktrace"]["type"],
                 "raw",
             )
             self.assertEqual(
-                capture_call[2]["$exception_list"][0]["stacktrace"]["frames"][0]["filename"],
+                capture_call[2]["$exception_list"][0]["stacktrace"]["frames"][0][
+                    "filename"
+                ],
                 "posthog/test/test_client.py",
             )
             self.assertEqual(
-                capture_call[2]["$exception_list"][0]["stacktrace"]["frames"][0]["function"],
+                capture_call[2]["$exception_list"][0]["stacktrace"]["frames"][0][
+                    "function"
+                ],
                 "test_basic_capture_exception_with_no_exception_given",
             )
             self.assertEqual(
-                capture_call[2]["$exception_list"][0]["stacktrace"]["frames"][0]["module"], "posthog.test.test_client"
+                capture_call[2]["$exception_list"][0]["stacktrace"]["frames"][0][
+                    "module"
+                ],
+                "posthog.test.test_client",
+            )
+            self.assertEqual(
+                capture_call[2]["$exception_list"][0]["stacktrace"]["frames"][0][
+                    "in_app"
+                ],
+                True,
             )
-            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:
@@ -267,16 +293,26 @@ class TestClient(unittest.TestCase):
     def test_capture_exception_logs_when_enabled(self):
         client = Client(FAKE_TEST_API_KEY, log_captured_exceptions=True)
         with self.assertLogs("posthog", level="ERROR") as logs:
-            client.capture_exception(Exception("test exception"), "distinct_id", path="one/two/three")
-            self.assertEqual(logs.output[0], "ERROR:posthog:test exception\nNoneType: None")
+            client.capture_exception(
+                Exception("test exception"), "distinct_id", path="one/two/three"
+            )
+            self.assertEqual(
+                logs.output[0], "ERROR:posthog:test exception\nNoneType: None"
+            )
             self.assertEqual(getattr(logs.records[0], "path"), "one/two/three")
 
     @mock.patch("posthog.client.flags")
     def test_basic_capture_with_feature_flags(self, patch_flags):
         patch_flags.return_value = {"featureFlags": {"beta-feature": "random-variant"}}
 
-        client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, personal_api_key=FAKE_TEST_API_KEY)
-        success, msg = client.capture("distinct_id", "python test event", send_feature_flags=True)
+        client = Client(
+            FAKE_TEST_API_KEY,
+            on_error=self.set_fail,
+            personal_api_key=FAKE_TEST_API_KEY,
+        )
+        success, msg = client.capture(
+            "distinct_id", "python test event", send_feature_flags=True
+        )
         client.flush()
         self.assertTrue(success)
         self.assertFalse(self.failed)
@@ -295,7 +331,11 @@ class TestClient(unittest.TestCase):
     @mock.patch("posthog.client.flags")
     def test_basic_capture_with_locally_evaluated_feature_flags(self, patch_flags):
         patch_flags.return_value = {"featureFlags": {"beta-feature": "random-variant"}}
-        client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, personal_api_key=FAKE_TEST_API_KEY)
+        client = Client(
+            FAKE_TEST_API_KEY,
+            on_error=self.set_fail,
+            personal_api_key=FAKE_TEST_API_KEY,
+        )
 
         multivariate_flag = {
             "id": 1,
@@ -307,7 +347,12 @@ class TestClient(unittest.TestCase):
                 "groups": [
                     {
                         "properties": [
-                            {"key": "email", "type": "person", "value": "test@posthog.com", "operator": "exact"}
+                            {
+                                "key": "email",
+                                "type": "person",
+                                "value": "test@posthog.com",
+                                "operator": "exact",
+                            }
                         ],
                         "rollout_percentage": 100,
                     },
@@ -317,12 +362,27 @@ class TestClient(unittest.TestCase):
                 ],
                 "multivariate": {
                     "variants": [
-                        {"key": "first-variant", "name": "First Variant", "rollout_percentage": 50},
-                        {"key": "second-variant", "name": "Second Variant", "rollout_percentage": 25},
-                        {"key": "third-variant", "name": "Third Variant", "rollout_percentage": 25},
+                        {
+                            "key": "first-variant",
+                            "name": "First Variant",
+                            "rollout_percentage": 50,
+                        },
+                        {
+                            "key": "second-variant",
+                            "name": "Second Variant",
+                            "rollout_percentage": 25,
+                        },
+                        {
+                            "key": "third-variant",
+                            "name": "Third Variant",
+                            "rollout_percentage": 25,
+                        },
                     ]
                 },
-                "payloads": {"first-variant": "some-payload", "third-variant": {"a": "json"}},
+                "payloads": {
+                    "first-variant": "some-payload",
+                    "third-variant": {"a": "json"},
+                },
             },
         }
         basic_flag = {
@@ -375,9 +435,13 @@ 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)
-        self.assertEqual(msg["properties"]["$feature/beta-feature-local"], "third-variant")
+        self.assertEqual(
+            msg["properties"]["$feature/beta-feature-local"], "third-variant"
+        )
         self.assertEqual(msg["properties"]["$feature/false-flag"], False)
-        self.assertEqual(msg["properties"]["$active_feature_flags"], ["beta-feature-local"])
+        self.assertEqual(
+            msg["properties"]["$active_feature_flags"], ["beta-feature-local"]
+        )
         assert "$feature/beta-feature" not in msg["properties"]
 
         self.assertEqual(patch_flags.call_count, 0)
@@ -415,7 +479,11 @@ class TestClient(unittest.TestCase):
     @mock.patch("posthog.client.flags")
     def test_dont_override_capture_with_local_flags(self, patch_flags):
         patch_flags.return_value = {"featureFlags": {"beta-feature": "random-variant"}}
-        client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, personal_api_key=FAKE_TEST_API_KEY)
+        client = Client(
+            FAKE_TEST_API_KEY,
+            on_error=self.set_fail,
+            personal_api_key=FAKE_TEST_API_KEY,
+        )
 
         multivariate_flag = {
             "id": 1,
@@ -427,7 +495,12 @@ class TestClient(unittest.TestCase):
                 "groups": [
                     {
                         "properties": [
-                            {"key": "email", "type": "person", "value": "test@posthog.com", "operator": "exact"}
+                            {
+                                "key": "email",
+                                "type": "person",
+                                "value": "test@posthog.com",
+                                "operator": "exact",
+                            }
                         ],
                         "rollout_percentage": 100,
                     },
@@ -437,12 +510,27 @@ class TestClient(unittest.TestCase):
                 ],
                 "multivariate": {
                     "variants": [
-                        {"key": "first-variant", "name": "First Variant", "rollout_percentage": 50},
-                        {"key": "second-variant", "name": "Second Variant", "rollout_percentage": 25},
-                        {"key": "third-variant", "name": "Third Variant", "rollout_percentage": 25},
+                        {
+                            "key": "first-variant",
+                            "name": "First Variant",
+                            "rollout_percentage": 50,
+                        },
+                        {
+                            "key": "second-variant",
+                            "name": "Second Variant",
+                            "rollout_percentage": 25,
+                        },
+                        {
+                            "key": "third-variant",
+                            "name": "Third Variant",
+                            "rollout_percentage": 25,
+                        },
                     ]
                 },
-                "payloads": {"first-variant": "some-payload", "third-variant": {"a": "json"}},
+                "payloads": {
+                    "first-variant": "some-payload",
+                    "third-variant": {"a": "json"},
+                },
             },
         }
         basic_flag = {
@@ -470,7 +558,9 @@ class TestClient(unittest.TestCase):
         client.feature_flags = [multivariate_flag, basic_flag]
 
         success, msg = client.capture(
-            "distinct_id", "python test event", {"$feature/beta-feature-local": "my-custom-variant"}
+            "distinct_id",
+            "python test event",
+            {"$feature/beta-feature-local": "my-custom-variant"},
         )
         client.flush()
         self.assertTrue(success)
@@ -482,8 +572,12 @@ 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)
-        self.assertEqual(msg["properties"]["$feature/beta-feature-local"], "my-custom-variant")
-        self.assertEqual(msg["properties"]["$active_feature_flags"], ["beta-feature-local"])
+        self.assertEqual(
+            msg["properties"]["$feature/beta-feature-local"], "my-custom-variant"
+        )
+        self.assertEqual(
+            msg["properties"]["$active_feature_flags"], ["beta-feature-local"]
+        )
         assert "$feature/beta-feature" not in msg["properties"]
         assert "$feature/person-flag" not in msg["properties"]
 
@@ -492,11 +586,21 @@ class TestClient(unittest.TestCase):
     @mock.patch("posthog.client.flags")
     def test_basic_capture_with_feature_flags_returns_active_only(self, patch_flags):
         patch_flags.return_value = {
-            "featureFlags": {"beta-feature": "random-variant", "alpha-feature": True, "off-feature": False}
+            "featureFlags": {
+                "beta-feature": "random-variant",
+                "alpha-feature": True,
+                "off-feature": False,
+            }
         }
 
-        client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, personal_api_key=FAKE_TEST_API_KEY)
-        success, msg = client.capture("distinct_id", "python test event", send_feature_flags=True)
+        client = Client(
+            FAKE_TEST_API_KEY,
+            on_error=self.set_fail,
+            personal_api_key=FAKE_TEST_API_KEY,
+        )
+        success, msg = client.capture(
+            "distinct_id", "python test event", send_feature_flags=True
+        )
         client.flush()
         self.assertTrue(success)
         self.assertFalse(self.failed)
@@ -510,7 +614,10 @@ class TestClient(unittest.TestCase):
         self.assertEqual(msg["properties"]["$lib_version"], VERSION)
         self.assertEqual(msg["properties"]["$feature/beta-feature"], "random-variant")
         self.assertEqual(msg["properties"]["$feature/alpha-feature"], True)
-        self.assertEqual(msg["properties"]["$active_feature_flags"], ["beta-feature", "alpha-feature"])
+        self.assertEqual(
+            msg["properties"]["$active_feature_flags"],
+            ["beta-feature", "alpha-feature"],
+        )
 
         self.assertEqual(patch_flags.call_count, 1)
         patch_flags.assert_called_with(
@@ -521,13 +628,19 @@ class TestClient(unittest.TestCase):
             groups={},
             person_properties=None,
             group_properties=None,
-            disable_geoip=True,
+            geoip_disable=True,
         )
 
     @mock.patch("posthog.client.flags")
-    def test_basic_capture_with_feature_flags_and_disable_geoip_returns_correctly(self, patch_flags):
+    def test_basic_capture_with_feature_flags_and_disable_geoip_returns_correctly(
+        self, patch_flags
+    ):
         patch_flags.return_value = {
-            "featureFlags": {"beta-feature": "random-variant", "alpha-feature": True, "off-feature": False}
+            "featureFlags": {
+                "beta-feature": "random-variant",
+                "alpha-feature": True,
+                "off-feature": False,
+            }
         }
 
         client = Client(
@@ -538,7 +651,12 @@ class TestClient(unittest.TestCase):
             disable_geoip=True,
             feature_flags_request_timeout_seconds=12,
         )
-        success, msg = client.capture("distinct_id", "python test event", send_feature_flags=True, disable_geoip=False)
+        success, msg = client.capture(
+            "distinct_id",
+            "python test event",
+            send_feature_flags=True,
+            disable_geoip=False,
+        )
         client.flush()
         self.assertTrue(success)
         self.assertFalse(self.failed)
@@ -552,7 +670,10 @@ class TestClient(unittest.TestCase):
         self.assertEqual(msg["properties"]["$lib_version"], VERSION)
         self.assertEqual(msg["properties"]["$feature/beta-feature"], "random-variant")
         self.assertEqual(msg["properties"]["$feature/alpha-feature"], True)
-        self.assertEqual(msg["properties"]["$active_feature_flags"], ["beta-feature", "alpha-feature"])
+        self.assertEqual(
+            msg["properties"]["$active_feature_flags"],
+            ["beta-feature", "alpha-feature"],
+        )
 
         self.assertEqual(patch_flags.call_count, 1)
         patch_flags.assert_called_with(
@@ -563,15 +684,23 @@ class TestClient(unittest.TestCase):
             groups={},
             person_properties=None,
             group_properties=None,
-            disable_geoip=False,
+            geoip_disable=False,
         )
 
     @mock.patch("posthog.client.flags")
-    def test_basic_capture_with_feature_flags_switched_off_doesnt_send_them(self, patch_flags):
+    def test_basic_capture_with_feature_flags_switched_off_doesnt_send_them(
+        self, patch_flags
+    ):
         patch_flags.return_value = {"featureFlags": {"beta-feature": "random-variant"}}
 
-        client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, personal_api_key=FAKE_TEST_API_KEY)
-        success, msg = client.capture("distinct_id", "python test event", send_feature_flags=False)
+        client = Client(
+            FAKE_TEST_API_KEY,
+            on_error=self.set_fail,
+            personal_api_key=FAKE_TEST_API_KEY,
+        )
+        success, msg = client.capture(
+            "distinct_id", "python test event", send_feature_flags=False
+        )
         client.flush()
         self.assertTrue(success)
         self.assertFalse(self.failed)
@@ -591,7 +720,9 @@ class TestClient(unittest.TestCase):
         # A large number that loses precision in node:
         # node -e "console.log(157963456373623802 + 1)" > 157963456373623800
         client = self.client
-        success, msg = client.capture(distinct_id=157963456373623802, event="python test event")
+        success, msg = client.capture(
+            distinct_id=157963456373623802, event="python test event"
+        )
         client.flush()
         self.assertTrue(success)
         self.assertFalse(self.failed)
@@ -627,7 +758,10 @@ class TestClient(unittest.TestCase):
         )
 
         self.assertTrue(success)
-        self.assertEqual(msg["properties"]["$groups"], {"company": "id:5", "instance": "app.posthog.com"})
+        self.assertEqual(
+            msg["properties"]["$groups"],
+            {"company": "id:5", "instance": "app.posthog.com"},
+        )
 
     def test_basic_identify(self):
         client = self.client
@@ -644,7 +778,10 @@ class TestClient(unittest.TestCase):
     def test_advanced_identify(self):
         client = self.client
         success, msg = client.identify(
-            "distinct_id", {"trait": "value"}, timestamp=datetime(2014, 9, 3), uuid="new-uuid"
+            "distinct_id",
+            {"trait": "value"},
+            timestamp=datetime(2014, 9, 3),
+            uuid="new-uuid",
         )
 
         self.assertTrue(success)
@@ -671,7 +808,12 @@ class TestClient(unittest.TestCase):
 
     def test_advanced_set(self):
         client = self.client
-        success, msg = client.set("distinct_id", {"trait": "value"}, timestamp=datetime(2014, 9, 3), uuid="new-uuid")
+        success, msg = client.set(
+            "distinct_id",
+            {"trait": "value"},
+            timestamp=datetime(2014, 9, 3),
+            uuid="new-uuid",
+        )
 
         self.assertTrue(success)
 
@@ -698,7 +840,10 @@ class TestClient(unittest.TestCase):
     def test_advanced_set_once(self):
         client = self.client
         success, msg = client.set_once(
-            "distinct_id", {"trait": "value"}, timestamp=datetime(2014, 9, 3), uuid="new-uuid"
+            "distinct_id",
+            {"trait": "value"},
+            timestamp=datetime(2014, 9, 3),
+            uuid="new-uuid",
         )
 
         self.assertTrue(success)
@@ -732,7 +877,9 @@ class TestClient(unittest.TestCase):
         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")
+        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")
@@ -752,7 +899,11 @@ class TestClient(unittest.TestCase):
 
     def test_advanced_group_identify(self):
         success, msg = self.client.group_identify(
-            "organization", "id:5", {"trait": "value"}, timestamp=datetime(2014, 9, 3), uuid="new-uuid"
+            "organization",
+            "id:5",
+            {"trait": "value"},
+            timestamp=datetime(2014, 9, 3),
+            uuid="new-uuid",
         )
 
         self.assertTrue(success)
@@ -814,7 +965,9 @@ class TestClient(unittest.TestCase):
         client.flush()
         self.assertTrue(success)
         self.assertEqual(msg["distinct_id"], "distinct_id")
-        self.assertEqual(msg["properties"]["$current_url"], "https://posthog.com/contact")
+        self.assertEqual(
+            msg["properties"]["$current_url"], "https://posthog.com/contact"
+        )
 
     def test_basic_page_distinct_uuid(self):
         client = self.client
@@ -824,7 +977,9 @@ class TestClient(unittest.TestCase):
         client.flush()
         self.assertTrue(success)
         self.assertEqual(msg["distinct_id"], str(distinct_id))
-        self.assertEqual(msg["properties"]["$current_url"], "https://posthog.com/contact")
+        self.assertEqual(
+            msg["properties"]["$current_url"], "https://posthog.com/contact"
+        )
 
     def test_advanced_page(self):
         client = self.client
@@ -839,7 +994,9 @@ class TestClient(unittest.TestCase):
         self.assertTrue(success)
 
         self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00")
-        self.assertEqual(msg["properties"]["$current_url"], "https://posthog.com/contact")
+        self.assertEqual(
+            msg["properties"]["$current_url"], "https://posthog.com/contact"
+        )
         self.assertEqual(msg["properties"]["property"], "value")
         self.assertEqual(msg["properties"]["$lib"], "posthog-python")
         self.assertEqual(msg["properties"]["$lib_version"], VERSION)
@@ -910,14 +1067,18 @@ class TestClient(unittest.TestCase):
         self.assertFalse(self.failed)
 
     def test_user_defined_flush_at(self):
-        client = Client(FAKE_TEST_API_KEY, on_error=self.fail, flush_at=10, flush_interval=3)
+        client = Client(
+            FAKE_TEST_API_KEY, on_error=self.fail, flush_at=10, flush_interval=3
+        )
 
         def mock_post_fn(*args, **kwargs):
             self.assertEqual(len(kwargs["batch"]), 10)
 
         # the post function should be called 2 times, with a batch size of 10
         # each time.
-        with mock.patch("posthog.consumer.batch_post", side_effect=mock_post_fn) as mock_post:
+        with mock.patch(
+            "posthog.consumer.batch_post", side_effect=mock_post_fn
+        ) as mock_post:
             for _ in range(20):
                 client.identify("distinct_id", {"trait": "value"})
             time.sleep(1)
@@ -998,11 +1159,15 @@ class TestClient(unittest.TestCase):
 
     def test_disable_geoip_override_on_events(self):
         client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, disable_geoip=False)
-        _, capture_msg = client.set("distinct_id", {"a": "b", "c": "d"}, disable_geoip=True)
+        _, capture_msg = client.set(
+            "distinct_id", {"a": "b", "c": "d"}, disable_geoip=True
+        )
         client.flush()
         self.assertEqual(capture_msg["properties"]["$geoip_disable"], True)
 
-        _, identify_msg = client.page("distinct_id", "http://a.com", {"trait": "value"}, disable_geoip=False)
+        _, identify_msg = client.page(
+            "distinct_id", "http://a.com", {"trait": "value"}, disable_geoip=False
+        )
         client.flush()
         self.assertEqual("$geoip_disable" not in identify_msg["properties"], True)
 
@@ -1015,7 +1180,11 @@ class TestClient(unittest.TestCase):
     @mock.patch("posthog.client.flags")
     def test_disable_geoip_default_on_decide(self, patch_flags):
         patch_flags.return_value = {
-            "featureFlags": {"beta-feature": "random-variant", "alpha-feature": True, "off-feature": False}
+            "featureFlags": {
+                "beta-feature": "random-variant",
+                "alpha-feature": True,
+                "off-feature": False,
+            }
         }
         client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, disable_geoip=False)
         client.get_feature_flag("random_key", "some_id", disable_geoip=True)
@@ -1027,10 +1196,12 @@ class TestClient(unittest.TestCase):
             groups={},
             person_properties={"distinct_id": "some_id"},
             group_properties={},
-            disable_geoip=True,
+            geoip_disable=True,
         )
         patch_flags.reset_mock()
-        client.feature_enabled("random_key", "feature_enabled_distinct_id", disable_geoip=True)
+        client.feature_enabled(
+            "random_key", "feature_enabled_distinct_id", disable_geoip=True
+        )
         patch_flags.assert_called_with(
             "random_key",
             "https://us.i.posthog.com",
@@ -1039,7 +1210,7 @@ class TestClient(unittest.TestCase):
             groups={},
             person_properties={"distinct_id": "feature_enabled_distinct_id"},
             group_properties={},
-            disable_geoip=True,
+            geoip_disable=True,
         )
         patch_flags.reset_mock()
         client.get_all_flags_and_payloads("all_flags_payloads_id")
@@ -1051,7 +1222,7 @@ class TestClient(unittest.TestCase):
             groups={},
             person_properties={"distinct_id": "all_flags_payloads_id"},
             group_properties={},
-            disable_geoip=False,
+            geoip_disable=False,
         )
 
     @mock.patch("posthog.client.Poller")
@@ -1069,9 +1240,18 @@ class TestClient(unittest.TestCase):
     @mock.patch("posthog.client.flags")
     def test_default_properties_get_added_properly(self, patch_flags):
         patch_flags.return_value = {
-            "featureFlags": {"beta-feature": "random-variant", "alpha-feature": True, "off-feature": False}
+            "featureFlags": {
+                "beta-feature": "random-variant",
+                "alpha-feature": True,
+                "off-feature": False,
+            }
         }
-        client = Client(FAKE_TEST_API_KEY, host="http://app2.posthog.com", on_error=self.set_fail, disable_geoip=False)
+        client = Client(
+            FAKE_TEST_API_KEY,
+            host="http://app2.posthog.com",
+            on_error=self.set_fail,
+            disable_geoip=False,
+        )
         client.get_feature_flag(
             "random_key",
             "some_id",
@@ -1090,7 +1270,7 @@ class TestClient(unittest.TestCase):
                 "company": {"$group_key": "id:5", "x": "y"},
                 "instance": {"$group_key": "app.posthog.com"},
             },
-            disable_geoip=False,
+            geoip_disable=False,
         )
 
         patch_flags.reset_mock()
@@ -1116,12 +1296,14 @@ class TestClient(unittest.TestCase):
                 "company": {"$group_key": "group_override"},
                 "instance": {"$group_key": "app.posthog.com"},
             },
-            disable_geoip=False,
+            geoip_disable=False,
         )
 
         patch_flags.reset_mock()
         # test nones
-        client.get_all_flags_and_payloads("some_id", groups={}, person_properties=None, group_properties=None)
+        client.get_all_flags_and_payloads(
+            "some_id", groups={}, person_properties=None, group_properties=None
+        )
         patch_flags.assert_called_with(
             "random_key",
             "http://app2.posthog.com",
@@ -1130,7 +1312,7 @@ class TestClient(unittest.TestCase):
             groups={},
             person_properties={"distinct_id": "some_id"},
             group_properties={},
-            disable_geoip=False,
+            geoip_disable=False,
         )
 
     @parameterized.expand(
@@ -1197,12 +1379,17 @@ class TestClient(unittest.TestCase):
 
                 # Set up platform-specific mocks
                 if platform_method:
-                    getattr(mock_platform, platform_method).return_value = platform_return
+                    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)):
+                    with mock.patch(
+                        "posthog.client.get_os_info",
+                        return_value=(expected_os, expected_os_version),
+                    ):
                         from posthog.client import system_context
 
                         context = system_context()
@@ -1225,7 +1412,11 @@ class TestClient(unittest.TestCase):
     @mock.patch("posthog.client.flags")
     def test_get_decide_returns_normalized_decide_response(self, patch_flags):
         patch_flags.return_value = {
-            "featureFlags": {"beta-feature": "random-variant", "alpha-feature": True, "off-feature": False},
+            "featureFlags": {
+                "beta-feature": "random-variant",
+                "alpha-feature": True,
+                "off-feature": False,
+            },
             "featureFlagPayloads": {"beta-feature": '{"some": "data"}'},
             "errorsWhileComputingFlags": False,
             "requestId": "test-id",
@@ -1292,10 +1483,15 @@ class TestClient(unittest.TestCase):
         client = Client(FAKE_TEST_API_KEY)
 
         # Test 100% rollout - should use flags
-        with mock.patch("posthog.client.is_token_in_rollout", return_value=True) as mock_rollout:
+        with mock.patch(
+            "posthog.client.is_token_in_rollout", return_value=True
+        ) as mock_rollout:
             client.get_flags_decision("distinct_id")
             mock_rollout.assert_called_with(
-                FAKE_TEST_API_KEY, 1, included_hashes=INCLUDED_HASHES, excluded_hashes=EXCLUDED_HASHES
+                FAKE_TEST_API_KEY,
+                1,
+                included_hashes=INCLUDED_HASHES,
+                excluded_hashes=EXCLUDED_HASHES,
             )
             patch_flags.assert_called_once()
             patch_decide.assert_not_called()
@@ -1307,25 +1503,39 @@ class TestClient(unittest.TestCase):
         included_hashes = {token_hash}
 
         # Should be included due to specific hash, even with 0% rollout
-        self.assertTrue(expr=is_token_in_rollout(token, percentage=0.0, included_hashes=included_hashes))
+        self.assertTrue(
+            expr=is_token_in_rollout(
+                token, percentage=0.0, included_hashes=included_hashes
+            )
+        )
 
         # Should not be included with 0% rollout and no specific hash
         self.assertFalse(is_token_in_rollout(token, percentage=0.0))
 
         # Should be included with 100% rollout regardless of specific hash
         self.assertTrue(is_token_in_rollout(token, percentage=1.0))
-        self.assertTrue(is_token_in_rollout(token, percentage=1.0, included_hashes=included_hashes))
+        self.assertTrue(
+            is_token_in_rollout(token, percentage=1.0, included_hashes=included_hashes)
+        )
 
         # Test deterministic behavior - same token should always give same result
         hash_float = int(token_hash[:8], 16) / 0xFFFFFFFF
         percentage = hash_float + 0.1  # Just above the hash value
 
         self.assertTrue(is_token_in_rollout(token, percentage))
-        self.assertFalse(is_token_in_rollout(token, percentage - 0.2))  # Just below the hash value
+        self.assertFalse(
+            is_token_in_rollout(token, percentage - 0.2)
+        )  # Just below the hash value
 
         # Test that the token exclusion works correctly
-        self.assertFalse(is_token_in_rollout(token, percentage=1.0, excluded_hashes={token_hash}))
+        self.assertFalse(
+            is_token_in_rollout(token, percentage=1.0, excluded_hashes={token_hash})
+        )
 
         # Should work for other specific token hashes
         # Include our API key
-        self.assertTrue(is_token_in_rollout("sTMFPsFhdP1Ssg", percentage=0.1, included_hashes=INCLUDED_HASHES))
+        self.assertTrue(
+            is_token_in_rollout(
+                "sTMFPsFhdP1Ssg", percentage=0.1, included_hashes=INCLUDED_HASHES
+            )
+        )

posthog/test/test_consumer.py

@@ -58,7 +58,11 @@ class TestConsumer(unittest.TestCase):
         with mock.patch("posthog.consumer.batch_post") as mock_post:
             consumer.start()
             for i in range(0, 3):
-                track = {"type": "track", "event": "python event %d" % i, "distinct_id": "distinct_id"}
+                track = {
+                    "type": "track",
+                    "event": "python event %d" % i,
+                    "distinct_id": "distinct_id",
+                }
                 q.put(track)
                 time.sleep(flush_interval * 1.1)
             self.assertEqual(mock_post.call_count, 3)
@@ -69,11 +73,17 @@ class TestConsumer(unittest.TestCase):
         q = Queue()
         flush_interval = 0.5
         flush_at = 10
-        consumer = Consumer(q, TEST_API_KEY, flush_at=flush_at, flush_interval=flush_interval)
+        consumer = Consumer(
+            q, TEST_API_KEY, flush_at=flush_at, flush_interval=flush_interval
+        )
         with mock.patch("posthog.consumer.batch_post") as mock_post:
             consumer.start()
             for i in range(0, flush_at * 2):
-                track = {"type": "track", "event": "python event %d" % i, "distinct_id": "distinct_id"}
+                track = {
+                    "type": "track",
+                    "event": "python event %d" % i,
+                    "distinct_id": "distinct_id",
+                }
                 q.put(track)
             time.sleep(flush_interval * 1.1)
             self.assertEqual(mock_post.call_count, 2)
@@ -91,8 +101,14 @@ class TestConsumer(unittest.TestCase):
 
         mock_post.call_count = 0
 
-        with mock.patch("posthog.consumer.batch_post", mock.Mock(side_effect=mock_post)):
-            track = {"type": "track", "event": "python event", "distinct_id": "distinct_id"}
+        with mock.patch(
+            "posthog.consumer.batch_post", mock.Mock(side_effect=mock_post)
+        ):
+            track = {
+                "type": "track",
+                "event": "python event",
+                "distinct_id": "distinct_id",
+            }
             # request() should succeed if the number of exceptions raised is
             # less than the retries paramater.
             if exception_count <= consumer.retries:
@@ -107,7 +123,8 @@ class TestConsumer(unittest.TestCase):
                     self.assertEqual(exc, expected_exception)
                 else:
                     self.fail(
-                        "request() should raise an exception if still failing after %d retries" % consumer.retries
+                        "request() should raise an exception if still failing after %d retries"
+                        % consumer.retries
                     )
 
     def test_request_retry(self):
@@ -148,7 +165,12 @@ class TestConsumer(unittest.TestCase):
         properties = {}
         for n in range(0, 500):
             properties[str(n)] = "one_long_property_value_to_build_a_big_event"
-        track = {"type": "track", "event": "python event", "distinct_id": "distinct_id", "properties": properties}
+        track = {
+            "type": "track",
+            "event": "python event",
+            "distinct_id": "distinct_id",
+            "properties": properties,
+        }
         msg_size = len(json.dumps(track).encode())
         # Let's capture 8MB of data to trigger two batches
         n_msgs = int(8_000_000 / msg_size)
@@ -158,10 +180,15 @@ class TestConsumer(unittest.TestCase):
             res.status_code = 200
             request_size = len(data.encode())
             # Batches close after the first message bringing it bigger than BATCH_SIZE_LIMIT, let's add 10% of margin
-            self.assertTrue(request_size < (5 * 1024 * 1024) * 1.1, "batch size (%d) higher than limit" % request_size)
+            self.assertTrue(
+                request_size < (5 * 1024 * 1024) * 1.1,
+                "batch size (%d) higher than limit" % request_size,
+            )
             return res
 
-        with mock.patch("posthog.request._session.post", side_effect=mock_post_fn) as mock_post:
+        with mock.patch(
+            "posthog.request._session.post", side_effect=mock_post_fn
+        ) as mock_post:
             consumer.start()
             for _ in range(0, n_msgs + 2):
                 q.put(track)

posthog/test/test_feature_flag.py

@@ -10,8 +10,17 @@ class TestFeatureFlag(unittest.TestCase):
             "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"},
+            "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)
@@ -20,10 +29,21 @@ class TestFeatureFlag(unittest.TestCase):
         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")
+            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")
+            flag.metadata,
+            FlagMetadata(
+                id=1,
+                payload='{"some": "json"}',
+                version=2,
+                description="test-description",
+            ),
         )
 
     def test_feature_flag_from_json_minimal(self):
@@ -44,7 +64,11 @@ class TestFeatureFlag(unittest.TestCase):
             "key": "test-flag",
             "enabled": True,
             "variant": "test-variant",
-            "reason": {"code": "matched_condition", "condition_index": 0, "description": "Matched condition set 1"},
+            "reason": {
+                "code": "matched_condition",
+                "condition_index": 0,
+                "description": "Matched condition set 1",
+            },
         }
 
         flag = FeatureFlag.from_json(resp)
@@ -53,13 +77,22 @@ class TestFeatureFlag(unittest.TestCase):
         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")
+            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'"}
+        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)
@@ -77,7 +110,12 @@ class TestFeatureFlag(unittest.TestCase):
 
     def test_flag_metadata_from_json(self):
         # Test with complete data
-        resp = {"id": 123, "payload": {"key": "value"}, "version": 1, "description": "Test flag"}
+        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"})
@@ -106,7 +144,12 @@ class TestFeatureFlag(unittest.TestCase):
                 "condition_index": 1,
                 "description": "User is in segment 'beta_users'",
             },
-            "metadata": {"id": 123, "payload": {"key": "value"}, "version": 1, "description": "Test flag"},
+            "metadata": {
+                "id": 123,
+                "payload": {"key": "value"},
+                "version": 1,
+                "description": "Test flag",
+            },
         }
         flag = FeatureFlag.from_json(resp)
         self.assertEqual(flag.key, "test-flag")
@@ -131,7 +174,11 @@ class TestFeatureFlag(unittest.TestCase):
 
     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"}}
+        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)

posthog/test/test_feature_flag_result.py

@@ -9,7 +9,9 @@ from posthog.types import FeatureFlag, FeatureFlagResult, FlagMetadata, FlagReas
 
 class TestFeatureFlagResult(unittest.TestCase):
     def test_from_bool_value_and_payload(self):
-        result = FeatureFlagResult.from_value_and_payload("test-flag", True, "[1, 2, 3]")
+        result = FeatureFlagResult.from_value_and_payload(
+            "test-flag", True, "[1, 2, 3]"
+        )
 
         self.assertEqual(result.key, "test-flag")
         self.assertEqual(result.enabled, True)
@@ -17,7 +19,9 @@ class TestFeatureFlagResult(unittest.TestCase):
         self.assertEqual(result.payload, [1, 2, 3])
 
     def test_from_false_value_and_payload(self):
-        result = FeatureFlagResult.from_value_and_payload("test-flag", False, '{"some": "value"}')
+        result = FeatureFlagResult.from_value_and_payload(
+            "test-flag", False, '{"some": "value"}'
+        )
 
         self.assertEqual(result.key, "test-flag")
         self.assertEqual(result.enabled, False)
@@ -25,7 +29,9 @@ class TestFeatureFlagResult(unittest.TestCase):
         self.assertEqual(result.payload, {"some": "value"})
 
     def test_from_variant_value_and_payload(self):
-        result = FeatureFlagResult.from_value_and_payload("test-flag", "control", "true")
+        result = FeatureFlagResult.from_value_and_payload(
+            "test-flag", "control", "true"
+        )
 
         self.assertEqual(result.key, "test-flag")
         self.assertEqual(result.enabled, True)
@@ -33,7 +39,9 @@ class TestFeatureFlagResult(unittest.TestCase):
         self.assertEqual(result.payload, True)
 
     def test_from_none_value_and_payload(self):
-        result = FeatureFlagResult.from_value_and_payload("test-flag", None, '{"some": "value"}')
+        result = FeatureFlagResult.from_value_and_payload(
+            "test-flag", None, '{"some": "value"}'
+        )
         self.assertIsNone(result)
 
     def test_from_boolean_flag_details(self):
@@ -41,8 +49,12 @@ class TestFeatureFlagResult(unittest.TestCase):
             key="test-flag",
             enabled=True,
             variant=None,
-            metadata=FlagMetadata(id=1, version=1, description="test-flag", payload='"Some string"'),
-            reason=FlagReason(code="test-reason", description="test-reason", condition_index=0),
+            metadata=FlagMetadata(
+                id=1, version=1, description="test-flag", payload='"Some string"'
+            ),
+            reason=FlagReason(
+                code="test-reason", description="test-reason", condition_index=0
+            ),
         )
 
         result = FeatureFlagResult.from_flag_details(flag_details)
@@ -57,11 +69,17 @@ class TestFeatureFlagResult(unittest.TestCase):
             key="test-flag",
             enabled=True,
             variant=None,
-            metadata=FlagMetadata(id=1, version=1, description="test-flag", payload='"Some string"'),
-            reason=FlagReason(code="test-reason", description="test-reason", condition_index=0),
+            metadata=FlagMetadata(
+                id=1, version=1, description="test-flag", payload='"Some string"'
+            ),
+            reason=FlagReason(
+                code="test-reason", description="test-reason", condition_index=0
+            ),
         )
 
-        result = FeatureFlagResult.from_flag_details(flag_details, override_match_value="control")
+        result = FeatureFlagResult.from_flag_details(
+            flag_details, override_match_value="control"
+        )
 
         self.assertEqual(result.key, "test-flag")
         self.assertEqual(result.enabled, True)
@@ -73,11 +91,17 @@ class TestFeatureFlagResult(unittest.TestCase):
             key="test-flag",
             enabled=True,
             variant="control",
-            metadata=FlagMetadata(id=1, version=1, description="test-flag", payload='{"some": "value"}'),
-            reason=FlagReason(code="test-reason", description="test-reason", condition_index=0),
+            metadata=FlagMetadata(
+                id=1, version=1, description="test-flag", payload='{"some": "value"}'
+            ),
+            reason=FlagReason(
+                code="test-reason", description="test-reason", condition_index=0
+            ),
         )
 
-        result = FeatureFlagResult.from_flag_details(flag_details, override_match_value=True)
+        result = FeatureFlagResult.from_flag_details(
+            flag_details, override_match_value=True
+        )
 
         self.assertEqual(result.key, "test-flag")
         self.assertEqual(result.enabled, True)
@@ -89,11 +113,17 @@ class TestFeatureFlagResult(unittest.TestCase):
             key="test-flag",
             enabled=True,
             variant="control",
-            metadata=FlagMetadata(id=1, version=1, description="test-flag", payload='{"some": "value"}'),
-            reason=FlagReason(code="test-reason", description="test-reason", condition_index=0),
+            metadata=FlagMetadata(
+                id=1, version=1, description="test-flag", payload='{"some": "value"}'
+            ),
+            reason=FlagReason(
+                code="test-reason", description="test-reason", condition_index=0
+            ),
         )
 
-        result = FeatureFlagResult.from_flag_details(flag_details, override_match_value=False)
+        result = FeatureFlagResult.from_flag_details(
+            flag_details, override_match_value=False
+        )
 
         self.assertEqual(result.key, "test-flag")
         self.assertEqual(result.enabled, False)
@@ -105,8 +135,12 @@ class TestFeatureFlagResult(unittest.TestCase):
             key="test-flag",
             enabled=True,
             variant="control",
-            metadata=FlagMetadata(id=1, version=1, description="test-flag", payload='{"some": "value"}'),
-            reason=FlagReason(code="test-reason", description="test-reason", condition_index=0),
+            metadata=FlagMetadata(
+                id=1, version=1, description="test-flag", payload='{"some": "value"}'
+            ),
+            reason=FlagReason(
+                code="test-reason", description="test-reason", condition_index=0
+            ),
         )
 
         result = FeatureFlagResult.from_flag_details(flag_details)
@@ -126,8 +160,12 @@ class TestFeatureFlagResult(unittest.TestCase):
             key="test-flag",
             enabled=True,
             variant=None,
-            metadata=FlagMetadata(id=1, version=1, description="test-flag", payload=None),
-            reason=FlagReason(code="test-reason", description="test-reason", condition_index=0),
+            metadata=FlagMetadata(
+                id=1, version=1, description="test-flag", payload=None
+            ),
+            reason=FlagReason(
+                code="test-reason", description="test-reason", condition_index=0
+            ),
         )
 
         result = FeatureFlagResult.from_flag_details(flag_details)
@@ -300,7 +338,9 @@ class TestGetFeatureFlagResult(unittest.TestCase):
             },
         }
 
-        flag_result = self.client.get_feature_flag_result("person-flag", "some-distinct-id")
+        flag_result = self.client.get_feature_flag_result(
+            "person-flag", "some-distinct-id"
+        )
         self.assertEqual(flag_result.enabled, True)
         self.assertEqual(flag_result.variant, None)
         self.assertEqual(flag_result.payload, 300)
@@ -385,7 +425,9 @@ class TestGetFeatureFlagResult(unittest.TestCase):
             },
         }
 
-        flag_result = self.client.get_feature_flag_result("no-person-flag", "some-distinct-id")
+        flag_result = self.client.get_feature_flag_result(
+            "no-person-flag", "some-distinct-id"
+        )
 
         self.assertIsNone(flag_result)
         patch_capture.assert_called_with(

posthog/test/test_module.py

@@ -15,7 +15,9 @@ class TestModule(unittest.TestCase):
 
     def setUp(self):
         self.failed = False
-        self.posthog = Posthog("testsecret", host="http://localhost:8000", on_error=self.failed)
+        self.posthog = Posthog(
+            "testsecret", host="http://localhost:8000", on_error=self.failed
+        )
 
     def test_no_api_key(self):
         self.posthog.api_key = None

posthog/test/test_request.py

@@ -6,20 +6,35 @@ import mock
 import pytest
 import requests
 
-from posthog.request import DatetimeSerializer, QuotaLimitError, batch_post, decide, determine_server_host
+from posthog.request import (
+    DatetimeSerializer,
+    QuotaLimitError,
+    batch_post,
+    decide,
+    determine_server_host,
+)
 from posthog.test.test_utils import TEST_API_KEY
 
 
 class TestRequests(unittest.TestCase):
     def test_valid_request(self):
-        res = batch_post(TEST_API_KEY, batch=[{"distinct_id": "distinct_id", "event": "python event", "type": "track"}])
+        res = batch_post(
+            TEST_API_KEY,
+            batch=[
+                {"distinct_id": "distinct_id", "event": "python event", "type": "track"}
+            ],
+        )
         self.assertEqual(res.status_code, 200)
 
     def test_invalid_request_error(self):
-        self.assertRaises(Exception, batch_post, "testsecret", "https://t.posthog.com", False, "[{]")
+        self.assertRaises(
+            Exception, batch_post, "testsecret", "https://t.posthog.com", False, "[{]"
+        )
 
     def test_invalid_host(self):
-        self.assertRaises(Exception, batch_post, "testsecret", "t.posthog.com/", batch=[])
+        self.assertRaises(
+            Exception, batch_post, "testsecret", "t.posthog.com/", batch=[]
+        )
 
     def test_datetime_serialization(self):
         data = {"created": datetime(2012, 3, 4, 5, 6, 7, 891011)}
@@ -35,14 +50,26 @@ class TestRequests(unittest.TestCase):
 
     def test_should_not_timeout(self):
         res = batch_post(
-            TEST_API_KEY, batch=[{"distinct_id": "distinct_id", "event": "python event", "type": "track"}], timeout=15
+            TEST_API_KEY,
+            batch=[
+                {"distinct_id": "distinct_id", "event": "python event", "type": "track"}
+            ],
+            timeout=15,
         )
         self.assertEqual(res.status_code, 200)
 
     def test_should_timeout(self):
         with self.assertRaises(requests.ReadTimeout):
             batch_post(
-                "key", batch=[{"distinct_id": "distinct_id", "event": "python event", "type": "track"}], timeout=0.0001
+                "key",
+                batch=[
+                    {
+                        "distinct_id": "distinct_id",
+                        "event": "python event",
+                        "type": "track",
+                    }
+                ],
+                timeout=0.0001,
             )
 
     def test_quota_limited_response(self):
@@ -68,7 +95,11 @@ class TestRequests(unittest.TestCase):
         mock_response = requests.Response()
         mock_response.status_code = 200
         mock_response._content = json.dumps(
-            {"featureFlags": {"flag1": True}, "featureFlagPayloads": {}, "errorsWhileComputingFlags": False}
+            {
+                "featureFlags": {"flag1": True},
+                "featureFlagPayloads": {},
+                "errorsWhileComputingFlags": False,
+            }
         ).encode("utf-8")
 
         with mock.patch("posthog.request._session.post", return_value=mock_response):

posthog/test/test_scopes.py

@@ -0,0 +1,138 @@
+import unittest
+from unittest.mock import patch
+
+from posthog.scopes import clear_tags, get_tags, new_context, scoped, tag
+
+
+class TestScopes(unittest.TestCase):
+    def setUp(self):
+        # Reset any context between tests
+        clear_tags()
+
+    def test_tag_and_get_tags(self):
+        tag("key1", "value1")
+        tag("key2", 2)
+
+        tags = get_tags()
+        assert tags["key1"] == "value1"
+        assert tags["key2"] == 2
+
+    def test_clear_tags(self):
+        tag("key1", "value1")
+        assert get_tags()["key1"] == "value1"
+
+        clear_tags()
+        assert get_tags() == {}
+
+    def test_new_context_isolation(self):
+        # Set tag in outer context
+        tag("outer", "value")
+
+        with new_context(fresh=True):
+            # Inner context should start empty
+            assert get_tags() == {}
+
+            # Set tag in inner context
+            tag("inner", "value")
+            assert get_tags()["inner"] == "value"
+
+            # Outer tag should not be visible
+            self.assertNotIn("outer", get_tags())
+
+        with new_context(fresh=False):
+            # Inner context should start empty
+            assert get_tags() == {"outer": "value"}
+
+        # After exiting context, inner tag should be gone
+        self.assertNotIn("inner", get_tags())
+
+        # Outer tag should still be there
+        assert get_tags()["outer"] == "value"
+
+    def test_nested_contexts(self):
+        tag("level1", "value1")
+
+        with new_context(fresh=True):
+            tag("level2", "value2")
+
+            with new_context(fresh=True):
+                tag("level3", "value3")
+                assert get_tags() == {"level3": "value3"}
+
+            # Back to level 2
+            assert get_tags() == {"level2": "value2"}
+
+        # Back to level 1
+        assert get_tags() == {"level1": "value1"}
+
+    @patch("posthog.capture_exception")
+    def test_scoped_decorator_success(self, mock_capture):
+        @scoped()
+        def successful_function(x, y):
+            tag("x", x)
+            tag("y", y)
+            return x + y
+
+        result = successful_function(1, 2)
+
+        # Function should execute normally
+        assert result == 3
+
+        # No exception should be captured
+        mock_capture.assert_not_called()
+
+        # Context should be cleared after function execution
+        assert get_tags() == {}
+
+    @patch("posthog.capture_exception")
+    def test_scoped_decorator_exception(self, mock_capture):
+        test_exception = ValueError("Test exception")
+
+        def check_context_on_capture(exception, **kwargs):
+            # Assert tags are available when capture_exception is called
+            current_tags = get_tags()
+            assert current_tags.get("important_context") == "value"
+
+        mock_capture.side_effect = check_context_on_capture
+
+        @scoped()
+        def failing_function():
+            tag("important_context", "value")
+            raise test_exception
+
+        # Function should raise the exception
+        with self.assertRaises(ValueError):
+            failing_function()
+
+        # Verify capture_exception was called
+        mock_capture.assert_called_once_with(test_exception)
+
+        # Context should be cleared after function execution
+        assert get_tags() == {}
+
+    @patch("posthog.capture_exception")
+    def test_new_context_exception_handling(self, mock_capture):
+        test_exception = RuntimeError("Context exception")
+
+        def check_context_on_capture(exception, **kwargs):
+            # Assert inner context tags are available when capture_exception is called
+            current_tags = get_tags()
+            assert current_tags.get("inner_context") == "inner_value"
+
+        mock_capture.side_effect = check_context_on_capture
+
+        # Set up outer context
+        tag("outer_context", "outer_value")
+
+        try:
+            with new_context():
+                tag("inner_context", "inner_value")
+                raise test_exception
+        except RuntimeError:
+            pass  # Expected exception
+
+        # Verify capture_exception was called
+        mock_capture.assert_called_once_with(test_exception)
+
+        # Outer context should still be intact
+        assert get_tags()["outer_context"] == "outer_value"

posthog/test/test_size_limited_dict.py

@@ -0,0 +1,24 @@
+import unittest
+
+from parameterized import parameterized
+
+from posthog import utils
+
+
+class TestSizeLimitedDict(unittest.TestCase):
+    @parameterized.expand([(10, 100), (5, 20), (20, 200)])
+    def test_size_limited_dict(self, size: int, iterations: int) -> None:
+        values = utils.SizeLimitedDict(size, lambda _: -1)
+
+        for i in range(iterations):
+            values[i] = i
+
+            assert values[i] == i
+            assert len(values) == i % size + 1
+
+            if i % size == 0:
+                # old numbers should've been removed
+                self.assertIsNone(values.get(i - 1))
+                self.assertIsNone(values.get(i - 3))
+                self.assertIsNone(values.get(i - 5))
+                self.assertIsNone(values.get(i - 9))

posthog/test/test_types.py

@@ -22,9 +22,16 @@ class TestTypes(unittest.TestCase):
                     enabled=True,
                     variant="test-variant",
                     reason=FlagReason(
-                        code="matched_condition", condition_index=0, description="Matched condition set 1"
+                        code="matched_condition",
+                        condition_index=0,
+                        description="Matched condition set 1",
+                    ),
+                    metadata=FlagMetadata(
+                        id=1,
+                        payload='{"some": "json"}',
+                        version=2,
+                        description="test-description",
                     ),
-                    metadata=FlagMetadata(id=1, payload='{"some": "json"}', version=2, description="test-description"),
                 )
             },
             "errorsWhileComputingFlags": has_errors,
@@ -39,10 +46,21 @@ class TestTypes(unittest.TestCase):
         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")
+            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")
+            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")
@@ -64,7 +82,9 @@ class TestTypes(unittest.TestCase):
         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.assertEqual(
+            flag.metadata, LegacyFlagMetadata(payload='{"some": "json-payload"}')
+        )
         self.assertFalse(result["errorsWhileComputingFlags"])
         self.assertEqual(result["requestId"], "test-id")
         # Verify legacy fields are removed
@@ -99,18 +119,29 @@ class TestTypes(unittest.TestCase):
                     enabled=True,
                     variant="test-variant",
                     reason=FlagReason(
-                        code="matched_condition", condition_index=0, description="Matched condition set 1"
+                        code="matched_condition",
+                        condition_index=0,
+                        description="Matched condition set 1",
+                    ),
+                    metadata=FlagMetadata(
+                        id=1,
+                        payload='{"some": "json"}',
+                        version=2,
+                        description="test-description",
                     ),
-                    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"
+                        code="matched_condition",
+                        condition_index=0,
+                        description="Matched condition set 1",
+                    ),
+                    metadata=FlagMetadata(
+                        id=1, payload=None, version=2, description="test-description"
                     ),
-                    metadata=FlagMetadata(id=1, payload=None, version=2, description="test-description"),
                 ),
                 "disabled-flag": FeatureFlag(
                     key="disabled-flag",
@@ -129,7 +160,9 @@ class TestTypes(unittest.TestCase):
         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.assertEqual(
+            result["featureFlagPayloads"]["my-variant-flag"], '{"some": "json"}'
+        )
         self.assertNotIn("my-boolean-flag", result["featureFlagPayloads"])
         self.assertNotIn("disabled-flag", result["featureFlagPayloads"])
 

posthog/test/test_utils.py

@@ -7,6 +7,7 @@ from uuid import UUID
 
 import six
 from dateutil.tz import tzutc
+from parameterized import parameterized
 from pydantic import BaseModel
 from pydantic.v1 import BaseModel as BaseModelV1
 
@@ -17,17 +18,29 @@ FAKE_TEST_API_KEY = "random_key"
 
 
 class TestUtils(unittest.TestCase):
+    @parameterized.expand(
+        [
+            ("naive datetime should be naive", True),
+            ("timezone-aware datetime should not be naive", False),
+        ]
+    )
+    def test_is_naive(self, _name: str, expected_naive: bool):
+        if expected_naive:
+            dt = datetime.now()  # naive datetime
+        else:
+            dt = datetime.now(tz=tzutc())  # timezone-aware datetime
+
+        assert utils.is_naive(dt) is expected_naive
+
     def test_timezone_utils(self):
         now = datetime.now()
         utcnow = datetime.now(tz=tzutc())
-        self.assertTrue(utils.is_naive(now))
-        self.assertFalse(utils.is_naive(utcnow))
 
         fixed = utils.guess_timezone(now)
-        self.assertFalse(utils.is_naive(fixed))
+        assert utils.is_naive(fixed) is False
 
         shouldnt_be_edited = utils.guess_timezone(utcnow)
-        self.assertEqual(utcnow, shouldnt_be_edited)
+        assert utcnow == shouldnt_be_edited
 
     def test_clean(self):
         simple = {
@@ -54,12 +67,12 @@ class TestUtils(unittest.TestCase):
         pre_clean_keys = combined.keys()
 
         utils.clean(combined)
-        self.assertEqual(combined.keys(), pre_clean_keys)
+        assert 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",
+        assert (
+            utils.clean(UUID("12345678123456781234567812345678"))
+            == "12345678-1234-5678-1234-567812345678"
         )
 
     def test_clean_with_dates(self):
@@ -67,26 +80,27 @@ class TestUtils(unittest.TestCase):
             "birthdate": date(1980, 1, 1),
             "registration": datetime.now(tz=tzutc()),
         }
-        self.assertEqual(dict_with_dates, utils.clean(dict_with_dates))
+        assert dict_with_dates == utils.clean(dict_with_dates)
 
     def test_bytes(self):
-        if six.PY3:
-            item = bytes(10)
-        else:
-            item = bytearray(10)
-
+        item = bytes(10)
         utils.clean(item)
+        assert utils.clean(item) == "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
 
     def test_clean_fn(self):
         cleaned = utils.clean({"fn": lambda x: x, "number": 4})
-        self.assertEqual(cleaned["number"], 4)
-        # TODO: fixme, different behavior on python 2 and 3
-        if "fn" in cleaned:
-            self.assertEqual(cleaned["fn"], None)
+        assert cleaned == {"fn": None, "number": 4}
 
-    def test_remove_slash(self):
-        self.assertEqual("http://posthog.io", utils.remove_trailing_slash("http://posthog.io/"))
-        self.assertEqual("http://posthog.io", utils.remove_trailing_slash("http://posthog.io"))
+    @parameterized.expand(
+        [
+            ("http://posthog.io/", "http://posthog.io"),
+            ("http://posthog.io", "http://posthog.io"),
+            ("https://example.com/path/", "https://example.com/path"),
+            ("https://example.com/path", "https://example.com/path"),
+        ]
+    )
+    def test_remove_slash(self, input_url, expected_url):
+        assert expected_url == utils.remove_trailing_slash(input_url)
 
     def test_clean_pydantic(self):
         class ModelV2(BaseModel):
@@ -101,19 +115,26 @@ class TestUtils(unittest.TestCase):
         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"}},
-        )
+        assert utils.clean(ModelV2(foo="1", bar=2)) == {
+            "foo": "1",
+            "bar": 2,
+            "baz": None,
+        }
+        assert utils.clean(ModelV1(foo=1, bar="2")) == {"foo": 1, "bar": "2"}
+        assert utils.clean(NestedModel(foo=ModelV2(foo="1", bar=2, baz="3"))) == {
+            "foo": {"foo": "1", "bar": 2, "baz": "3"}
+        }
 
+    def test_clean_pydantic_like_class(self) -> None:
         class Dummy:
-            def model_dump(self, required_param):
-                pass
+            def model_dump(self, required_param: str) -> dict:
+                return {}
 
-        # Skips a class with a defined non-Pydantic `model_dump` method.
-        self.assertEqual(utils.clean({"test": Dummy()}), {})
+        # previously python 2 code would cause an error while cleaning,
+        # and this entire object would be None, and we would log an error
+        # let's allow ourselves to clean `Dummy` as None,
+        # without blatting the `test` key
+        assert utils.clean({"test": Dummy()}) == {"test": None}
 
     def test_clean_dataclass(self):
         @dataclass
@@ -130,47 +151,25 @@ class TestUtils(unittest.TestCase):
             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,
-                },
+        assert 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):
-        size = 10
-        values = utils.SizeLimitedDict(size, lambda _: -1)
-
-        for i in range(100):
-            values[i] = i
-
-            self.assertEqual(values[i], i)
-            self.assertEqual(len(values), i % size + 1)
-
-            if i % size == 0:
-                # old numbers should've been removed
-                self.assertIsNone(values.get(i - 1))
-                self.assertIsNone(values.get(i - 3))
-                self.assertIsNone(values.get(i - 5))
-                self.assertIsNone(values.get(i - 9))
+        }

posthog/types.py

@@ -78,7 +78,9 @@ class FeatureFlag:
         )
 
     @classmethod
-    def from_value_and_payload(cls, key: str, value: FlagValue, payload: Any) -> "FeatureFlag":
+    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,
@@ -160,7 +162,9 @@ class FeatureFlagResult:
 
     @classmethod
     def from_flag_details(
-        cls, details: Union[FeatureFlag, None], override_match_value: Optional[FlagValue] = None
+        cls,
+        details: Union[FeatureFlag, None],
+        override_match_value: Optional[FlagValue] = None,
     ) -> "FeatureFlagResult | None":
         """
         Create a FeatureFlagResult from a FeatureFlag object.
@@ -179,7 +183,9 @@ class FeatureFlagResult:
 
         if override_match_value is not None:
             enabled, variant = (
-                (True, override_match_value) if isinstance(override_match_value, str) else (override_match_value, None)
+                (True, override_match_value)
+                if isinstance(override_match_value, str)
+                else (override_match_value, None)
             )
         else:
             enabled, variant = (details.enabled, details.variant)
@@ -226,7 +232,9 @@ def normalize_flags_response(resp: Any) -> FlagsResponse:
         # 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))
+            flags[key] = FeatureFlag.from_value_and_payload(
+                key, value, featureFlagPayloads.get(key, None)
+            )
         resp["flags"] = flags
     return cast(FlagsResponse, resp)
 
@@ -252,7 +260,11 @@ def to_values(response: FlagsResponse) -> Optional[dict[str, FlagValue]]:
         return None
 
     flags = response.get("flags", {})
-    return {key: value.get_value() for key, value in flags.items() if isinstance(value, FeatureFlag)}
+    return {
+        key: value.get_value()
+        for key, value in flags.items()
+        if isinstance(value, FeatureFlag)
+    }
 
 
 def to_payloads(response: FlagsResponse) -> Optional[dict[str, str]]:

posthog/utils.py

@@ -5,6 +5,7 @@ from collections import defaultdict
 from dataclasses import asdict, is_dataclass
 from datetime import date, datetime, timezone
 from decimal import Decimal
+from typing import Any, Optional
 from uuid import UUID
 
 import six
@@ -52,7 +53,9 @@ def clean(item):
         return float(item)
     if isinstance(item, UUID):
         return str(item)
-    if 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
     if isinstance(item, (set, list, tuple)):
         return _clean_list(item)
@@ -99,14 +102,35 @@ def _clean_dataclass(dataclass_):
     return data
 
 
-def _coerce_unicode(cmplx):
+def _coerce_unicode(cmplx: Any) -> Optional[str]:
+    """
+    In theory, this method is only called
+    after many isinstance checks are carried out in `utils.clean`.
+    When we supported Python 2 it was safe to call `decode` on a `str`
+    but in Python 3 that will throw.
+    So, we check if the input is bytes and only call `decode` in that case.
+
+    Previously we would always call `decode` on the input
+    That would throw an error.
+    Then we would call `decode` on the stringified error
+    That would throw an error.
+    And then we would return `None`
+
+    To avoid a breaking change, we can maintain the behavior
+    that anything which did not have `decode` in Python 2
+    returns None.
+    """
+    item = None
     try:
-        item = cmplx.decode("utf-8", "strict")
-    except AttributeError as exception:
-        item = ":".join(exception)
-        item.decode("utf-8", "strict")
+        if isinstance(cmplx, bytes):
+            item = cmplx.decode("utf-8", "strict")
+        elif isinstance(cmplx, str):
+            item = cmplx
+    except Exception as exception:
+        item = ":".join(map(str, exception.args))
         log.warning("Error decoding: %s", item)
         return None
+
     return item
 
 

posthog/version.py

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

pyproject.toml

@@ -1,10 +1,93 @@
-[tool.black]
-line-length = 120
+[build-system]
+requires = ["setuptools>=61.0", "wheel"]
+build-backend = "setuptools.build_meta"
 
-[tool.isort]
-multi_line_output = 3
-include_trailing_comma = true
-force_grid_wrap = 8
-ensure_newline_before_comments = true
-line_length = 120
-virtual_env = "env"
+[project]
+name = "posthog"
+dynamic = ["version"]
+description = "Integrate PostHog into any python application."
+authors = [{ name = "PostHog", email = "hey@posthog.com" }]
+maintainers = [{ name = "PostHog", email = "hey@posthog.com" }]
+license = { text = "MIT License" }
+readme = "README.md"
+requires-python = ">=3.9"
+classifiers = [
+    "Development Status :: 5 - Production/Stable",
+    "Intended Audience :: Developers",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python",
+    "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",
+]
+dependencies = [
+    "requests>=2.7,<3.0",
+    "six>=1.5",
+    "python-dateutil>=2.2",
+    "backoff>=1.10.0",
+    "distro>=1.5.0",
+]
+
+[project.urls]
+Homepage = "https://github.com/posthog/posthog-python"
+Repository = "https://github.com/posthog/posthog-python"
+
+[project.optional-dependencies]
+dev = [
+    "django-stubs",
+    "lxml",
+    "mypy",
+    "mypy-baseline",
+    "types-mock",
+    "types-python-dateutil",
+    "types-requests",
+    "types-setuptools",
+    "types-six",
+    "pre-commit",
+    "pydantic",
+    "ruff",
+]
+test = [
+    "mock>=2.0.0",
+    "freezegun==1.5.1",
+    "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",
+    "google-genai",
+    "pydantic",
+    "parameterized>=0.8.1",
+]
+sentry = ["sentry-sdk", "django"]
+langchain = ["langchain>=0.2.0"]
+
+[tool.setuptools]
+packages = [
+    "posthog",
+    "posthog.ai",
+    "posthog.ai.langchain",
+    "posthog.ai.openai",
+    "posthog.ai.anthropic",
+    "posthog.ai.gemini",
+    "posthog.test",
+    "posthog.sentry",
+    "posthog.exception_integrations",
+]
+
+license-files = []
+
+[tool.setuptools.dynamic]
+version = { attr = "posthog.version.VERSION" }
+
+[tool.pytest.ini_options]
+asyncio_mode = "auto"
+asyncio_default_fixture_loop_scope = "function"

sentry_django_example/manage.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python
 """Django's command-line utility for administrative tasks."""
+
 import os
 import sys
 

sentry_django_example/sentry_django_example/settings.py

@@ -30,10 +30,12 @@ ALLOWED_HOSTS = []
 
 
 # PostHog Setup (can be a separate app)
-import posthog
+import posthog  # noqa: E402
 
 # You can find this key on the /setup page in PostHog
-posthog.api_key = "LXP6nQXvo-2TCqGVrWvPah8uJIyVykoMmhnEkEBi5PA"  # TODO: replace with your api key
+posthog.api_key = (
+    "LXP6nQXvo-2TCqGVrWvPah8uJIyVykoMmhnEkEBi5PA"  # TODO: replace with your api key
+)
 
 posthog.personal_api_key = ""
 
@@ -41,7 +43,7 @@ posthog.personal_api_key = ""
 # You can remove this line if you're using posthog.com
 posthog.host = "http://127.0.0.1:8000"
 
-from posthog.sentry.posthog_integration import PostHogIntegration
+from posthog.sentry.posthog_integration import PostHogIntegration  # noqa: E402
 
 PostHogIntegration.organization = "posthog"  # TODO: your sentry organization
 # PostHogIntegration.prefix = # TODO: your self hosted Sentry url. (default: https://sentry.io/organizations/)
@@ -50,8 +52,8 @@ PostHogIntegration.organization = "posthog"  # TODO: your sentry organization
 # we work around this by setting static class variables beforehand
 
 # Sentry Setup
-import sentry_sdk
-from sentry_sdk.integrations.django import DjangoIntegration
+import sentry_sdk  # noqa: E402
+from sentry_sdk.integrations.django import DjangoIntegration  # noqa: E402
 
 sentry_sdk.init(
     dsn="https://27ac54f7f4cf484abf1335436b0c52e5@o344752.ingest.sentry.io/5624115",  # TODO: your Sentry DSN here
@@ -66,7 +68,9 @@ sentry_sdk.init(
 )
 
 POSTHOG_DJANGO = {
-    "distinct_id": lambda request: str(uuid4())  # TODO: your logic for generating unique ID, given the request object
+    "distinct_id": lambda request: str(
+        uuid4()
+    )  # TODO: your logic for generating unique ID, given the request object
 }
 
 # Application definition

setup.cfg

@@ -3,3 +3,10 @@ universal = 1
 
 [tool:pytest]
 asyncio_mode = auto
+asyncio_default_fixture_loop_scope = function
+
+[flake8]
+# ignore E501 for line length
+# ignore W503 for line break before binary operator
+ignore = E501,W503
+max-line-length = 120

setup.py

@@ -8,97 +8,29 @@ except ImportError:
 
 # Don't import analytics-python module here, since deps may not be installed
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), "posthog"))
-from version import VERSION
+from version import VERSION  # noqa: E402
 
 long_description = """
-PostHog is developer-friendly, self-hosted product analytics. posthog-python is the python package.
+PostHog is developer-friendly, self-hosted product analytics.
+posthog-python is the python package.
 
 This package requires Python 3.9 or higher.
 """
 
-install_requires = [
-    "requests>=2.7,<3.0",
-    "six>=1.5",
-    "python-dateutil>=2.2",
-    "backoff>=1.10.0",
-    "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",
-    ],
-    "sentry": ["sentry-sdk", "django"],
-    "langchain": ["langchain>=0.2.0"],
-}
-
+# Minimal setup.py for backward compatibility
+# Most configuration is now in pyproject.toml
 setup(
     name="posthog",
     version=VERSION,
+    # Basic fields for backward compatibility
     url="https://github.com/posthog/posthog-python",
     author="Posthog",
     author_email="hey@posthog.com",
     maintainer="PostHog",
     maintainer_email="hey@posthog.com",
     test_suite="posthog.test.all",
-    packages=[
-        "posthog",
-        "posthog.ai",
-        "posthog.ai.langchain",
-        "posthog.ai.openai",
-        "posthog.ai.anthropic",
-        "posthog.test",
-        "posthog.sentry",
-        "posthog.exception_integrations",
-    ],
     license="MIT License",
-    install_requires=install_requires,
-    extras_require=extras_require,
     description="Integrate PostHog into any python application.",
     long_description=long_description,
-    classifiers=[
-        "Development Status :: 5 - Production/Stable",
-        "Intended Audience :: Developers",
-        "License :: OSI Approved :: MIT License",
-        "Operating System :: OS Independent",
-        "Programming Language :: Python",
-        "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",
-    ],
+    # This will fallback to pyproject.toml for detailed configuration
 )

setup_analytics.py

@@ -6,63 +6,31 @@ try:
 except ImportError:
     from distutils.core import setup
 
-# Don't import module here, since deps may not be installed
-sys.path.insert(0, os.path.join(os.path.dirname(__file__), "posthoganalytics"))
-from version import VERSION
+# Don't import analytics-python module here, since deps may not be installed
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "posthog"))
+from version import VERSION  # noqa: E402
 
 long_description = """
-PostHog is developer-friendly, self-hosted product analytics. posthog-python is the python package.
+PostHog is developer-friendly, self-hosted product analytics.
+posthog-python is the python package.
 
 This package requires Python 3.9 or higher.
 """
 
-install_requires = [
-    "requests>=2.7,<3.0",
-    "six>=1.5",
-    "python-dateutil>=2.2",
-    "backoff>=1.10.0",
-    "distro>=1.5.0",  # Required for Linux OS detection in Python 3.9+
-]
-
-tests_require = ["mock>=2.0.0"]
-
+# Minimal setup.py for backward compatibility
+# Most configuration is now in pyproject.toml
 setup(
     name="posthoganalytics",
     version=VERSION,
+    # Basic fields for backward compatibility
     url="https://github.com/posthog/posthog-python",
     author="Posthog",
     author_email="hey@posthog.com",
     maintainer="PostHog",
     maintainer_email="hey@posthog.com",
-    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",
-    ],
+    test_suite="posthog.test.all",
     license="MIT License",
-    install_requires=install_requires,
-    tests_require=tests_require,
-    extras_require={
-        "sentry": ["sentry-sdk", "django"],
-    },
     description="Integrate PostHog into any python application.",
     long_description=long_description,
-    classifiers=[
-        "Development Status :: 5 - Production/Stable",
-        "Intended Audience :: Developers",
-        "License :: OSI Approved :: MIT License",
-        "Operating System :: OS Independent",
-        "Programming Language :: Python",
-        "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",
-    ],
+    # This will fallback to pyproject.toml for detailed configuration
 )

simulator.py

@@ -28,7 +28,9 @@ parser.add_argument("--anonymousId", help="the anonymous user id to send the eve
 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)")
 
-parser.add_argument("--name", help="name of the screen or page to send with the message")
+parser.add_argument(
+    "--name", help="name of the screen or page to send with the message"
+)
 
 parser.add_argument("--traits", help="the identify/group traits to send (JSON-encoded)")
 
@@ -94,7 +96,13 @@ ch = logging.StreamHandler()
 ch.setLevel(logging.DEBUG)
 log.addHandler(ch)
 
-switcher = {"capture": capture, "page": page, "identify": identify, "set_once": set_once, "set": set}
+switcher = {
+    "capture": capture,
+    "page": page,
+    "identify": identify,
+    "set_once": set_once,
+    "set": set,
+}
 
 func = switcher.get(options.type)
 if func: