remove 'import posthog' (#255)

* chore: more fiddling to get release working

* fix

* fix

.github/workflows/ci.yml

@@ -9,38 +9,29 @@ jobs:
         runs-on: ubuntu-latest
 
         steps:
-            - uses: actions/checkout@v2
+            - uses: actions/checkout@85e6279cec87321a52edac9c87bce653a07cf6c2
               with:
                   fetch-depth: 1
 
-            - name: Set up Python 3.13
-              uses: actions/setup-python@v2
+            - name: Set up Python 3.11
+              uses: actions/setup-python@8d9ed9ac5c53483de85588cdf95a591a75ab9f55
               with:
-                  python-version: "3.13"
+                  python-version: 3.11.11
 
-            - uses: actions/cache@v3
+            - name: Install uv
+              uses: astral-sh/setup-uv@0c5e2b8115b80b4c7c5ddf6ffdd634974642d182 # v5.4.1
               with:
-                  path: ~/.cache/pip
-                  key: ${{ runner.os }}-pip-${{ hashFiles('setup.py') }}
-                  restore-keys: |
-                      ${{ runner.os }}-pip-
+                  enable-cache: true
+                  pyproject-file: 'pyproject.toml'
 
             - name: Install dev dependencies
+              shell: bash
               run: |
-                  python -m pip install -e .[dev]
-              if: steps.cache.outputs.cache-hit != 'true'
+                  UV_PROJECT_ENVIRONMENT=$pythonLocation uv sync --extra dev
 
-            - 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,18 +45,25 @@ 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 }}
 
-            - name: Install requirements.txt dependencies with pip
+            - name: Install uv
+              uses: astral-sh/setup-uv@0c5e2b8115b80b4c7c5ddf6ffdd634974642d182 # v5.4.1
+              with:
+                  enable-cache: true
+                  pyproject-file: 'pyproject.toml'
+
+            - name: Install test dependencies
+              shell: bash
               run: |
-                  python -m pip install -e .[test]
+                  UV_PROJECT_ENVIRONMENT=$pythonLocation uv sync --extra test
 
             - name: Run posthog tests
               run: |

.github/workflows/release.yaml

@@ -1,38 +1,51 @@
-name: 'Release'
+name: "Release"
 
 on:
-  - workflow_dispatch
+  push:
+    branches:
+      - master
+    paths:
+      - "posthog/version.py"
+  workflow_dispatch:
 
 jobs:
-    release:
-        name: Publish release
-        runs-on: ubuntu-20.04
+  release:
+    name: Publish release
+    runs-on: ubuntu-latest
+    env:
+      TWINE_USERNAME: __token__
+      TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }}
+    steps:
+      - name: Checkout the repository
+        uses: actions/checkout@85e6279cec87321a52edac9c87bce653a07cf6c2
+        with:
+          fetch-depth: 0
+          token: ${{ secrets.POSTHOG_BOT_GITHUB_TOKEN }}
+
+      - name: Set up Python
+        uses: actions/setup-python@8d9ed9ac5c53483de85588cdf95a591a75ab9f55
+        with:
+          python-version: 3.11.11
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@0c5e2b8115b80b4c7c5ddf6ffdd634974642d182 # v5.4.1
+        with:
+            enable-cache: true
+            pyproject-file: 'pyproject.toml'
+            
+      - name: Detect version
+        run: echo "REPO_VERSION=$(python3 posthog/version.py)" >> $GITHUB_ENV
+
+      - name: Prepare for building release
+        run: uv sync --extra dev
+
+      - name: Push releases to PyPI
+        run: uv run make release && uv run make release_analytics
+
+      - name: Create GitHub release
+        uses: actions/create-release@0cb9c9b65d5d1901c1f53e5e66eaf4afd303e70e # v1
         env:
-            TWINE_USERNAME: __token__
-            TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }}
-        steps:
-            - name: Checkout the repository
-              uses: actions/checkout@v2
-              with:
-                  fetch-depth: 0
-                  token: ${{ secrets.POSTHOG_BOT_GITHUB_TOKEN }}
-
-            - name: Set up Python
-              uses: actions/setup-python@v2
-
-            - 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
-
-            - name: Push release to PyPI
-              run: make release && make release_analytics
-
-            - name: Create GitHub release
-              uses: actions/create-release@v1
-              env:
-                  GITHUB_TOKEN: ${{ secrets.POSTHOG_BOT_GITHUB_TOKEN }}
-              with:
-                  tag_name: v${{ env.REPO_VERSION }}
-                  release_name: ${{ env.REPO_VERSION }}
+          GITHUB_TOKEN: ${{ secrets.POSTHOG_BOT_GITHUB_TOKEN }}
+        with:
+          tag_name: v${{ env.REPO_VERSION }}
+          release_name: ${{ env.REPO_VERSION }}

.gitignore

@@ -17,3 +17,4 @@ posthog-analytics
 .coverage
 pyrightconfig.json
 .env
+.DS_Store

.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

BEFORE_SEND.md

@@ -0,0 +1,237 @@
+# Before Send Hook
+
+The `before_send` parameter allows you to modify or filter events before they are sent to PostHog. This is useful for:
+
+- **Privacy**: Removing or masking sensitive data (PII)
+- **Filtering**: Dropping unwanted events (test events, internal users, etc.)
+- **Enhancement**: Adding custom properties to all events
+- **Transformation**: Modifying event names or property formats
+
+## Basic Usage
+
+```python
+import posthog
+from typing import Optional, Dict, Any
+
+def my_before_send(event: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+    """
+    Process event before sending to PostHog.
+    
+    Args:
+        event: The event dictionary containing 'event', 'distinct_id', 'properties', etc.
+    
+    Returns:
+        Modified event dictionary to send, or None to drop the event
+    """
+    # Your processing logic here
+    return event
+
+# Initialize client with before_send hook
+client = posthog.Client(
+    api_key="your-project-api-key",
+    before_send=my_before_send
+)
+```
+
+## Common Use Cases
+
+### 1. Filter Out Events
+
+```python
+from typing import Optional, Any
+
+def filter_events_by_property_or_event_name(event: dict[str, Any]) -> Optional[dict[str, Any]]:
+    """Drop events from internal users or test environments."""
+    properties = event.get("properties", {})
+    
+    # Choose some property from your events
+    event_source = properties.get("event_source", "")
+    if event_source.endswith("internal"):
+        return None  # Drop the event
+    
+    # Filter out test events
+    if event.get("event") == "test_event":
+        return None
+    
+    return event
+```
+
+### 2. Remove/Mask PII Data
+
+```python
+from typing import Optional, Any
+
+def scrub_pii(event: dict[str, Any]) -> Optional[dict[str, Any]]:
+    """Remove or mask personally identifiable information."""
+    properties = event.get("properties", {})
+    
+    # Mask email but keep domain for analytics
+    if "email" in properties:
+        email = properties["email"]
+        if "@" in email:
+            domain = email.split("@")[1]
+            properties["email"] = f"***@{domain}"
+        else:
+            properties["email"] = "***"
+    
+    # Remove sensitive fields entirely
+    sensitive_fields = ["my_business_info", "secret_things"]
+    for field in sensitive_fields:
+        properties.pop(field, None)
+    
+    return event
+```
+
+### 3. Add Custom Properties
+
+```python
+from typing import Optional, Any
+
+from datetime import datetime
+from typing import Optional, Any
+
+def add_context(event: dict[str, Any]) -> Optional[dict[str, Any]]:
+    """Add custom properties to all events."""
+    if "properties" not in event:
+        event["properties"] = {}
+    
+    event["properties"].update({
+        "app_version": "2.1.0",
+        "environment": "production", 
+        "processed_at": datetime.now().isoformat()
+    })
+    
+    return event
+```
+
+### 4. Transform Event Names
+
+```python
+from typing import Optional, Any
+
+def normalize_event_names(event: dict[str, Any]) -> Optional[dict[str, Any]]:
+    """Convert event names to a consistent format."""
+    original_event = event.get("event")
+    if original_event:
+        # Convert to snake_case
+        normalized = original_event.lower().replace(" ", "_").replace("-", "_")
+        event["event"] = f"app_{normalized}"
+    
+    return event
+```
+
+### 5. Log and drop in "dev" mode
+
+When running in local dev often, you want to log but drop all events
+
+
+```python
+from typing import Optional, Any
+
+def log_and_drop_all(event: dict[str, Any]) -> Optional[dict[str, Any]]:
+    """Convert event names to a consistent format."""
+    print(event)
+    
+    return None
+```
+
+### 6. Combined Processing
+
+```python
+from typing import Optional, Any
+
+def comprehensive_processor(event: dict[str, Any]) -> Optional[dict[str, Any]]:
+    """Apply multiple transformations in sequence."""
+    
+    # Step 1: Filter unwanted events
+    if should_drop_event(event):
+        return None
+    
+    # Step 2: Scrub PII
+    event = scrub_pii(event)
+    
+    # Step 3: Add context
+    event = add_context(event)
+    
+    # Step 4: Normalize names
+    event = normalize_event_names(event)
+    
+    return event
+
+def should_drop_event(event: dict[str, Any]) -> bool:
+    """Determine if event should be dropped."""
+    # Your filtering logic
+    return False
+```
+
+## Error Handling
+
+If your `before_send` function raises an exception, PostHog will:
+
+1. Log the error 
+2. Continue with the original, unmodified event
+3. Not crash your application
+
+```python
+from typing import Optional, Any
+
+def risky_before_send(event: dict[str, Any]) -> Optional[dict[str, Any]]:
+    # If this raises an exception, the original event will be sent
+    risky_operation()
+    return event
+```
+
+## Complete Example
+
+```python
+import posthog
+from typing import Optional, Any
+import re
+
+def production_before_send(event: dict[str, Any]) -> Optional[dict[str, Any]]:
+    try:
+        properties = event.get("properties", {})
+        
+        # 1. Filter out bot traffic
+        user_agent = properties.get("$user_agent", "")
+        if re.search(r'bot|crawler|spider', user_agent, re.I):
+            return None
+        
+        # 2. Filter out internal traffic  
+        ip = properties.get("$ip", "")
+        if ip.startswith("192.168.") or ip.startswith("10."):
+            return None
+        
+        # 3. Scrub email PII but keep domain
+        if "email" in properties:
+            email = properties["email"]
+            if "@" in email:
+                domain = email.split("@")[1]
+                properties["email"] = f"***@{domain}"
+        
+        # 4. Add custom context
+        properties.update({
+            "app_version": "1.0.0",
+            "build_number": "123"
+        })
+        
+        # 5. Normalize event name
+        if event.get("event"):
+            event["event"] = event["event"].lower().replace(" ", "_")
+        
+        return event
+        
+    except Exception as e:
+        # Log error but don't crash
+        print(f"Error in before_send: {e}")
+        return event  # Return original event on error
+
+# Usage
+client = posthog.Client(
+    api_key="your-api-key",
+    before_send=production_before_send
+)
+
+# All events will now be processed by your before_send function
+client.capture("user_123", "Page View", {"url": "/home"})
+```

CHANGELOG.md

@@ -1,3 +1,168 @@
+## 4.6.2 - 2025-06-09
+
+- fix: replace `import posthog` with direct method imports
+
+## 4.6.1 - 2025-06-09
+
+- fix: replace `import posthog` in `posthoganalytics` package
+
+## 4.6.0 - 2025-06-09
+
+- feat: add additional user and request context to captured exceptions via the Django integration
+- feat: Add `setup()` function to initialise default client
+
+## 4.5.0 - 2025-06-09
+
+- feat: add before_send callback (#249)
+
+## 4.4.2- 2025-06-09
+
+- empty point release to fix release automation
+
+## 4.4.1 2025-06-09
+
+- empty point release to fix release automation
+
+## 4.4.0 - 2025-06-09
+
+- Use the new `/flags` endpoint for all feature flag evaluations (don't fall back to `/decide` at all)
+
+## 4.3.2 - 2025-06-06
+
+1. 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
+
+2. fix: feature flag request use geoip_disable (#235)
+3. chore: pin actions versions (#210)
+4. fix: opinionated setup and clean fn fix (#240)
+5. 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
+2. Clarify Python 3.9+ is required
+
+## 4.0.0 - 2025-04-24
+
+1. Added new method `get_feature_flag_result` which returns a `FeatureFlagResult` object. This object breaks down the result of a feature flag into its enabled state, variant, and payload. The benefit of this method is it allows you to retrieve the result of a feature flag and its payload in a single API call. You can call `get_value` on the result to get the value of the feature flag, which is the same value returned by `get_feature_flag` (aka the string `variant` if the flag is a multivariate flag or the `boolean` value if the flag is a boolean flag).
+
+Example:
+
+```python
+result = posthog.get_feature_flag_result("my-flag", "distinct_id")
+print(result.enabled)     # True or False
+print(result.variant)     # 'the-variant-value' or None
+print(result.payload)     # {'foo': 'bar'}
+print(result.get_value()) # 'the-variant-value' or True or False
+print(result.reason)      # 'matched condition set 2' (Not available for local evaluation)
+```
+
+Breaking change:
+
+1. `get_feature_flag_payload` now deserializes payloads from JSON strings to `Any`. Previously, it returned the payload as a JSON encoded string.
+
+Before:
+
+```python
+payload = get_feature_flag_payload('key', 'distinct_id') # "{\"some\": \"payload\"}"
+```
+
+After:
+
+```python
+payload = get_feature_flag_payload('key', 'distinct_id') # {"some": "payload"}
+```
+
+## 3.25.0 – 2025-04-15
+
+1. Roll out new `/flags` endpoint to 100% of `/decide` traffic, excluding the top 10 customers.
+
+## 3.24.3 – 2025-04-15
+
+1. Fix hash inclusion/exclusion for flag rollout
+
+## 3.24.2 – 2025-04-15
+
+1. Roll out new /flags endpoint to 10% of /decide traffic
+
+## 3.24.1 – 2025-04-11
+
+1. Add `log_captured_exceptions` option to proxy setup
+
+## 3.24.0 – 2025-04-10
+
+1. Add config option to `log_captured_exceptions`
+
+## 3.23.0 – 2025-03-26
+
+1. Expand automatic retries to include read errors (e.g. RemoteDisconnected)
+
+## 3.22.0 – 2025-03-26
+
+1. Add more information to `$feature_flag_called` events.
+2. Support for the `/decide?v=4` endpoint which contains more information about feature flags.
+
+## 3.21.0 – 2025-03-17
+
+1. Support serializing dataclasses.
+
+## 3.20.0 – 2025-03-13
+
+1. Add support for OpenAI Responses API.
+
+## 3.19.2 – 2025-03-11
+
+1. Fix install requirements for analytics package
+
+## 3.19.1 – 2025-03-11
+
+1. Fix bug where None is sent as delta in azure
+
+## 3.19.0 – 2025-03-04
+
+1. Add support for tool calls in OpenAI and Anthropic.
+2. Add support for cached tokens.
+
+## 3.18.1 – 2025-03-03
+
+1. Improve quota-limited feature flag logs
+
+## 3.18.0 - 2025-02-28
+
+1. Add support for Azure OpenAI.
+
+## 3.17.0 - 2025-02-27
+
+1. The LangChain handler now captures tools in `$ai_generation` events, in property `$ai_tools`. This allows for displaying tools provided to the LLM call in PostHog UI. Note that support for `$ai_tools` in OpenAI and Anthropic SDKs is coming soon.
+
+## 3.16.0 - 2025-02-26
+
+1. feat: add some platform info to events (#198)
+
+## 3.15.1 - 2025-02-23
+
+1. Fix async client support for OpenAI.
+
+## 3.15.0 - 2025-02-19
+
+1. Support quota-limited feature flags
+
+## 3.14.2 - 2025-02-19
+
+1. Evaluate feature flag payloads with case sensitivity correctly. Fixes <https://github.com/PostHog/posthog-python/issues/178>
 
 ## 3.14.1 - 2025-02-18
 

LICENSE

@@ -20,3 +20,29 @@ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 SOFTWARE.
 
+---
+
+Some files in this codebase contain code from getsentry/sentry-javascript by Software, Inc. dba Sentry.
+In such cases it is explicitly stated in the file header. This license only applies to the relevant code in such cases.
+
+MIT License
+
+Copyright (c) 2012 Functional Software, Inc. dba Sentry
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of
+this software and associated documentation files (the "Software"), to deal in
+the Software without restriction, including without limitation the rights to
+use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
+of the Software, and to permit persons to whom the Software is furnished to do
+so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

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
@@ -17,16 +16,21 @@ release_analytics:
 	rm -rf posthoganalytics
 	mkdir posthoganalytics
 	cp -r posthog/* posthoganalytics/
-	find ./posthoganalytics -type f -exec sed -i '' -e 's/from posthog /from posthoganalytics /g' {} \;
-	find ./posthoganalytics -type f -exec sed -i '' -e 's/from posthog\./from posthoganalytics\./g' {} \;
+	find ./posthoganalytics -type f -name "*.py" -exec sed -i.bak -e 's/from posthog /from posthoganalytics /g' {} \;
+	find ./posthoganalytics -type f -name "*.py" -exec sed -i.bak -e 's/from posthog\./from posthoganalytics\./g' {} \;
+	find ./posthoganalytics -name "*.bak" -delete
 	rm -rf posthog
 	python setup_analytics.py sdist bdist_wheel
 	twine upload dist/*
 	mkdir posthog
-	find ./posthoganalytics -type f -exec sed -i '' -e 's/from posthoganalytics /from posthog /g' {} \;
-	find ./posthoganalytics -type f -exec sed -i '' -e 's/from posthoganalytics\./from posthog\./g' {} \;
+	find ./posthoganalytics -type f -name "*.py" -exec sed -i.bak -e 's/from posthoganalytics /from posthog /g' {} \;
+	find ./posthoganalytics -type f -name "*.py" -exec sed -i.bak -e 's/from posthoganalytics\./from posthog\./g' {} \;
+	find ./posthoganalytics -name "*.bak" -delete
 	cp -r posthoganalytics/* posthog/
 	rm -rf posthoganalytics
+	rm -f pyproject.toml
+	cp pyproject.toml.backup pyproject.toml
+	rm -f pyproject.toml.backup
 
 e2e_test:
 	.buildscripts/e2e.sh

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.
 
@@ -9,12 +16,29 @@ Please see the [Python integration docs](https://posthog.com/docs/integrations/p
 
 ### Testing Locally
 
-1. Run `python3 -m venv env` (creates virtual environment called "env")
+We recommend using [uv](https://docs.astral.sh/uv/). It's super fast.
+
+1. Run `uv venv env` (creates virtual environment called "env")
+    * or `python3 -m venv env`
 2. Run `source env/bin/activate` (activates the virtual environment)
-3. Run `python3 -m pip install -e ".[test]"` (installs the package in develop mode, along with test dependencies)
-4. Run `make test`
+3. Run `uv sync --extra dev --extra test` (installs the package in develop mode, along with test dependencies)
+    * or `pip install -e ".[dev,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 sync --extra dev --extra 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.
@@ -39,8 +63,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`.

bin/build

@@ -0,0 +1,8 @@
+#!/usr/bin/env bash
+#/ Usage: bin/build
+#/ Description: Runs linter and mypy
+source bin/helpers/_utils.sh
+set_source_and_root_dir
+
+flake8 posthog --ignore E501,W503
+mypy --no-site-packages --config-file mypy.ini . | mypy-baseline filter

bin/fmt

@@ -0,0 +1,14 @@
+#!/usr/bin/env bash
+#/ Usage: bin/fmt
+#/ Description: Formats and lints the code
+source bin/helpers/_utils.sh
+set_source_and_root_dir
+ensure_virtual_env
+
+if [[ "$1" == "--check" ]]; then
+  black --check .
+  isort --check-only .
+else
+  black .
+  isort .
+fi

bin/helpers/_utils.sh

@@ -0,0 +1,26 @@
+error() {
+    echo "$@" >&2
+}
+
+fatal() {
+    error "$@"
+    exit 1
+}
+
+set_source_and_root_dir() {
+    { set +x; } 2>/dev/null
+    source_dir="$( cd -P "$( dirname "$0" )" >/dev/null 2>&1 && pwd )"
+    root_dir=$(cd "$source_dir" && cd ../ && pwd)
+    cd "$root_dir"
+}
+
+ensure_virtual_env() {
+    if [ -z "$VIRTUAL_ENV" ]; then
+        echo "Virtual environment not activated. Activating now..."
+        if [ ! -f env/bin/activate ]; then
+            echo "Virtual environment not found. Please run 'python -m venv env' first."
+            exit 1
+        fi
+        source env/bin/activate
+    fi
+}

bin/setup

@@ -0,0 +1,12 @@
+#!/usr/bin/env bash
+#/ Usage: bin/setup
+#/ Description: Sets up the dependencies needed to develop this project
+source bin/helpers/_utils.sh
+set_source_and_root_dir
+
+if [ ! -d "env" ]; then
+  python3 -m venv env
+fi
+
+source env/bin/activate
+pip install -e ".[dev,test]"

bin/test

@@ -0,0 +1,10 @@
+#!/usr/bin/env bash
+#/ Usage: bin/test
+#/ Description: Runs all the unit tests for this project
+source bin/helpers/_utils.sh
+set_source_and_root_dir
+
+ensure_virtual_env
+
+# Pass through all arguments to pytest
+pytest "$@"

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

llm_observability_examples.py

@@ -1,220 +0,0 @@
-import os
-import uuid
-
-from pydantic import BaseModel
-
-import posthog
-from posthog.ai.openai import AsyncOpenAI, OpenAI
-
-# Example credentials - replace these with your own or use environment variables
-posthog.project_api_key = os.getenv("POSTHOG_PROJECT_API_KEY", "your-project-api-key")
-posthog.host = os.getenv("POSTHOG_HOST", "http://localhost:8000")  # Or https://app.posthog.com
-posthog.debug = True
-# change this to False to see usage events
-# posthog.privacy_mode = True
-
-openai_client = OpenAI(
-    api_key=os.getenv("OPENAI_API_KEY", "your-openai-api-key"),
-    posthog_client=posthog,
-)
-
-async_openai_client = AsyncOpenAI(
-    api_key=os.getenv("OPENAI_API_KEY", "your-openai-api-key"),
-    posthog_client=posthog,
-)
-
-
-def main_sync():
-    trace_id = str(uuid.uuid4())
-    print("Trace ID:", trace_id)
-    distinct_id = "test2_distinct_id"
-    properties = {"test_property": "test_value"}
-    groups = {"company": "test_company"}
-
-    try:
-        # basic_openai_call(distinct_id, trace_id, properties, groups)
-        # streaming_openai_call(distinct_id, trace_id, properties, groups)
-        # embedding_openai_call(distinct_id, trace_id, properties, groups)
-        # image_openai_call()
-        beta_openai_call(distinct_id, trace_id, properties, groups)
-    except Exception as e:
-        print("Error during OpenAI call:", str(e))
-
-
-async def main_async():
-    trace_id = str(uuid.uuid4())
-    print("Trace ID:", trace_id)
-    distinct_id = "test_distinct_id"
-    properties = {"test_property": "test_value"}
-    groups = {"company": "test_company"}
-
-    try:
-        await basic_async_openai_call(distinct_id, trace_id, properties, groups)
-        await streaming_async_openai_call(distinct_id, trace_id, properties, groups)
-        await embedding_async_openai_call(distinct_id, trace_id, properties, groups)
-        await image_async_openai_call()
-    except Exception as e:
-        print("Error during OpenAI call:", str(e))
-
-
-def basic_openai_call(distinct_id, trace_id, properties, groups):
-    response = openai_client.chat.completions.create(
-        model="gpt-4o-mini",
-        messages=[
-            {"role": "system", "content": "You are a complex problem solver."},
-            {"role": "user", "content": "Explain quantum computing in simple terms."},
-        ],
-        max_tokens=100,
-        temperature=0.7,
-        posthog_distinct_id=distinct_id,
-        posthog_trace_id=trace_id,
-        posthog_properties=properties,
-        posthog_groups=groups,
-    )
-    print(response)
-    if response and response.choices:
-        print("OpenAI response:", response.choices[0].message.content)
-    else:
-        print("No response or unexpected format returned.")
-    return response
-
-
-async def basic_async_openai_call(distinct_id, trace_id, properties, groups):
-    response = await async_openai_client.chat.completions.create(
-        model="gpt-4o-mini",
-        messages=[
-            {"role": "system", "content": "You are a complex problem solver."},
-            {"role": "user", "content": "Explain quantum computing in simple terms."},
-        ],
-        max_tokens=100,
-        temperature=0.7,
-        posthog_distinct_id=distinct_id,
-        posthog_trace_id=trace_id,
-        posthog_properties=properties,
-        posthog_groups=groups,
-    )
-    if response and hasattr(response, "choices"):
-        print("OpenAI response:", response.choices[0].message.content)
-    else:
-        print("No response or unexpected format returned.")
-    return response
-
-
-def streaming_openai_call(distinct_id, trace_id, properties, groups):
-
-    response = openai_client.chat.completions.create(
-        model="gpt-4o-mini",
-        messages=[
-            {"role": "system", "content": "You are a complex problem solver."},
-            {"role": "user", "content": "Explain quantum computing in simple terms."},
-        ],
-        max_tokens=100,
-        temperature=0.7,
-        stream=True,
-        posthog_distinct_id=distinct_id,
-        posthog_trace_id=trace_id,
-        posthog_properties=properties,
-        posthog_groups=groups,
-    )
-
-    for chunk in response:
-        if hasattr(chunk, "choices") and chunk.choices and len(chunk.choices) > 0:
-            print(chunk.choices[0].delta.content or "", end="")
-
-    return response
-
-
-async def streaming_async_openai_call(distinct_id, trace_id, properties, groups):
-    response = await async_openai_client.chat.completions.create(
-        model="gpt-4o-mini",
-        messages=[
-            {"role": "system", "content": "You are a complex problem solver."},
-            {"role": "user", "content": "Explain quantum computing in simple terms."},
-        ],
-        max_tokens=100,
-        temperature=0.7,
-        stream=True,
-        posthog_distinct_id=distinct_id,
-        posthog_trace_id=trace_id,
-        posthog_properties=properties,
-        posthog_groups=groups,
-    )
-
-    async for chunk in response:
-        if hasattr(chunk, "choices") and chunk.choices and len(chunk.choices) > 0:
-            print(chunk.choices[0].delta.content or "", end="")
-
-    return response
-
-
-# none instrumented
-def image_openai_call():
-    response = openai_client.images.generate(model="dall-e-3", prompt="A cute baby hedgehog", n=1, size="1024x1024")
-    print(response)
-    return response
-
-
-# none instrumented
-async def image_async_openai_call():
-    response = await async_openai_client.images.generate(
-        model="dall-e-3", prompt="A cute baby hedgehog", n=1, size="1024x1024"
-    )
-    print(response)
-    return response
-
-
-def embedding_openai_call(posthog_distinct_id, posthog_trace_id, posthog_properties, posthog_groups):
-    response = openai_client.embeddings.create(
-        input="The hedgehog is cute",
-        model="text-embedding-3-small",
-        posthog_distinct_id=posthog_distinct_id,
-        posthog_trace_id=posthog_trace_id,
-        posthog_properties=posthog_properties,
-        posthog_groups=posthog_groups,
-    )
-    print(response)
-    return response
-
-
-async def embedding_async_openai_call(posthog_distinct_id, posthog_trace_id, posthog_properties, posthog_groups):
-    response = await async_openai_client.embeddings.create(
-        input="The hedgehog is cute",
-        model="text-embedding-3-small",
-        posthog_distinct_id=posthog_distinct_id,
-        posthog_trace_id=posthog_trace_id,
-        posthog_properties=posthog_properties,
-        posthog_groups=posthog_groups,
-    )
-    print(response)
-    return response
-
-
-class CalendarEvent(BaseModel):
-    name: str
-    date: str
-    participants: list[str]
-
-
-def beta_openai_call(distinct_id, trace_id, properties, groups):
-    response = openai_client.beta.chat.completions.parse(
-        model="gpt-4o-mini",
-        messages=[
-            {"role": "system", "content": "Extract the event information."},
-            {"role": "user", "content": "Alice and Bob are going to a science fair on Friday."},
-        ],
-        response_format=CalendarEvent,
-        posthog_distinct_id=distinct_id,
-        posthog_trace_id=trace_id,
-        posthog_properties=properties,
-        posthog_groups=groups,
-    )
-    print(response)
-    return response
-
-
-# HOW TO RUN:
-# comment out one of these to run the other
-
-if __name__ == "__main__":
-    main_sync()
-# asyncio.run(main_async())

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]
@@ -35,27 +34,7 @@ posthog/sentry/posthog_integration.py:0: error: Statement is unreachable  [unrea
 posthog/ai/utils.py:0: error: Need type annotation for "output" (hint: "output: list[<type>] = ...")  [var-annotated]
 posthog/ai/utils.py:0: error: Function "builtins.any" is not valid as a type  [valid-type]
 posthog/ai/utils.py:0: note: Perhaps you meant "typing.Any" instead of "any"?
-posthog/ai/utils.py:0: error: Incompatible types in assignment (expression has type "UUID", variable has type "str | None")  [assignment]
 posthog/ai/utils.py:0: error: Function "builtins.any" is not valid as a type  [valid-type]
 posthog/ai/utils.py:0: note: Perhaps you meant "typing.Any" instead of "any"?
-posthog/ai/utils.py:0: error: Incompatible types in assignment (expression has type "UUID", variable has type "str | None")  [assignment]
 sentry_django_example/sentry_django_example/settings.py:0: error: Need type annotation for "ALLOWED_HOSTS" (hint: "ALLOWED_HOSTS: list[<type>] = ...")  [var-annotated]
 sentry_django_example/sentry_django_example/settings.py:0: error: Incompatible types in assignment (expression has type "str", variable has type "None")  [assignment]
-posthog/ai/openai/openai_async.py:0: error: Incompatible types in assignment (expression has type "UUID", variable has type "str | None")  [assignment]
-posthog/ai/openai/openai_async.py:0: error: Incompatible types in assignment (expression has type "UUID", variable has type "str | None")  [assignment]
-posthog/ai/openai/openai_async.py:0: error: Unpacked dict entry 11 has incompatible type "dict[str, Any] | None"; expected "SupportsKeysAndGetItem[str, Any]"  [dict-item]
-posthog/ai/openai/openai_async.py:0: error: Incompatible types in assignment (expression has type "UUID", variable has type "str | None")  [assignment]
-posthog/ai/openai/openai_async.py:0: error: Unpacked dict entry 8 has incompatible type "dict[str, Any] | None"; expected "SupportsKeysAndGetItem[str, Any]"  [dict-item]
-posthog/ai/openai/openai.py:0: error: Incompatible types in assignment (expression has type "UUID", variable has type "str | None")  [assignment]
-posthog/ai/openai/openai.py:0: error: Incompatible types in assignment (expression has type "UUID", variable has type "str | None")  [assignment]
-posthog/ai/openai/openai.py:0: error: Unpacked dict entry 11 has incompatible type "dict[str, Any] | None"; expected "SupportsKeysAndGetItem[str, Any]"  [dict-item]
-posthog/ai/openai/openai.py:0: error: Incompatible types in assignment (expression has type "UUID", variable has type "str | None")  [assignment]
-posthog/ai/openai/openai.py:0: error: Unpacked dict entry 8 has incompatible type "dict[str, Any] | None"; expected "SupportsKeysAndGetItem[str, Any]"  [dict-item]
-posthog/ai/anthropic/anthropic_async.py:0: error: Incompatible types in assignment (expression has type "UUID", variable has type "str | None")  [assignment]
-posthog/ai/anthropic/anthropic_async.py:0: error: Incompatible types in assignment (expression has type "UUID", variable has type "str | None")  [assignment]
-posthog/ai/anthropic/anthropic_async.py:0: error: Incompatible types in assignment (expression has type "UUID", variable has type "str | None")  [assignment]
-posthog/ai/anthropic/anthropic.py:0: error: Incompatible types in assignment (expression has type "UUID", variable has type "str | None")  [assignment]
-posthog/ai/anthropic/anthropic.py:0: error: Incompatible types in assignment (expression has type "UUID", variable has type "str | None")  [assignment]
-posthog/ai/anthropic/anthropic.py:0: error: Incompatible types in assignment (expression has type "UUID", variable has type "str | None")  [assignment]
-llm_observability_examples.py:0: error: Argument "posthog_client" to "OpenAI" has incompatible type Module; expected "Client"  [arg-type]
-llm_observability_examples.py:0: error: Argument "posthog_client" to "AsyncOpenAI" has incompatible type Module; expected "Client"  [arg-type]

mypy.ini

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

posthog/__init__.py

@@ -4,10 +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]
@@ -25,6 +34,7 @@ super_properties = None  # type: Optional[Dict]
 # Currently alpha, use at your own risk
 enable_exception_autocapture = False  # type: bool
 exception_autocapture_integrations = []  # type: List[Integrations]
+log_captured_exceptions = False  # type: bool
 # Used to determine in app paths for exception autocapture. Defaults to the current working directory
 project_root = None  # type: Optional[str]
 # Used for our AI observability feature to not capture any prompt or output just usage + metadata
@@ -313,6 +323,7 @@ def capture_exception(
     timestamp=None,  # type: Optional[datetime.datetime]
     uuid=None,  # type: Optional[str]
     groups=None,  # type: Optional[Dict]
+    **kwargs,
 ):
     # type: (...) -> Tuple[bool, dict]
     """
@@ -326,6 +337,7 @@ def capture_exception(
     Optionally you can submit
     - `properties`, which can be a dict with any information you'd like to add
     - `groups`, which is a dict of group type -> group key mappings
+    - remaining `kwargs` will be logged if `log_captured_exceptions` is enabled
 
     For example:
     ```python
@@ -354,6 +366,7 @@ def capture_exception(
         timestamp=timestamp,
         uuid=uuid,
         groups=groups,
+        **kwargs,
     )
 
 
@@ -403,7 +416,7 @@ def get_feature_flag(
     only_evaluate_locally=False,  # type: bool
     send_feature_flag_events=True,  # type: bool
     disable_geoip=None,  # type: Optional[bool]
-):
+) -> Optional[FeatureFlag]:
     """
     Get feature flag variant for users. Used with experiments.
     Example:
@@ -446,7 +459,7 @@ def get_all_flags(
     group_properties={},  # type: dict
     only_evaluate_locally=False,  # type: bool
     disable_geoip=None,  # type: Optional[bool]
-):
+) -> Optional[dict[str, FeatureFlag]]:
     """
     Get all flags for a given user.
     Example:
@@ -477,7 +490,7 @@ def get_feature_flag_payload(
     only_evaluate_locally=False,
     send_feature_flag_events=True,
     disable_geoip=None,  # type: Optional[bool]
-):
+) -> Optional[str]:
     return _proxy(
         "get_feature_flag_payload",
         key=key,
@@ -519,7 +532,7 @@ def get_all_flags_and_payloads(
     group_properties={},
     only_evaluate_locally=False,
     disable_geoip=None,  # type: Optional[bool]
-):
+) -> FlagsAndPayloads:
     return _proxy(
         "get_all_flags_and_payloads",
         distinct_id=distinct_id,
@@ -567,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(
@@ -589,6 +601,7 @@ def _proxy(method, *args, **kwargs):
             # This kind of initialisation is very annoying for exception capture. We need to figure out a way around this,
             # or deprecate this proxy option fully (it's already in the process of deprecation, no new clients should be using this method since like 5-6 months)
             enable_exception_autocapture=enable_exception_autocapture,
+            log_captured_exceptions=log_captured_exceptions,
             exception_autocapture_integrations=exception_autocapture_integrations,
         )
 
@@ -596,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
 
 
@@ -54,7 +61,7 @@ class WrappedMessages(Messages):
             **kwargs: Arguments passed to Anthropic's messages.create
         """
         if posthog_trace_id is None:
-            posthog_trace_id = uuid.uuid4()
+            posthog_trace_id = str(uuid.uuid4())
 
         if kwargs.get("stream", False):
             return self._create_streaming(
@@ -89,7 +96,7 @@ class WrappedMessages(Messages):
         **kwargs: Any,
     ):
         if posthog_trace_id is None:
-            posthog_trace_id = uuid.uuid4()
+            posthog_trace_id = str(uuid.uuid4())
 
         return self._create_streaming(
             posthog_distinct_id,
@@ -116,7 +123,7 @@ class WrappedMessages(Messages):
 
         def generator():
             nonlocal usage_stats
-            nonlocal accumulated_content
+            nonlocal accumulated_content  # noqa: F824
             try:
                 for event in response:
                     if hasattr(event, "usage") and event.usage:
@@ -125,6 +132,8 @@ class WrappedMessages(Messages):
                             for k in [
                                 "input_tokens",
                                 "output_tokens",
+                                "cache_read_input_tokens",
+                                "cache_creation_input_tokens",
                             ]
                         }
 
@@ -165,7 +174,7 @@ class WrappedMessages(Messages):
         output: str,
     ):
         if posthog_trace_id is None:
-            posthog_trace_id = uuid.uuid4()
+            posthog_trace_id = str(uuid.uuid4())
 
         event_properties = {
             "$ai_provider": "anthropic",
@@ -184,6 +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_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
 
 
@@ -54,7 +61,7 @@ class AsyncWrappedMessages(AsyncMessages):
             **kwargs: Arguments passed to Anthropic's messages.create
         """
         if posthog_trace_id is None:
-            posthog_trace_id = uuid.uuid4()
+            posthog_trace_id = str(uuid.uuid4())
 
         if kwargs.get("stream", False):
             return await self._create_streaming(
@@ -89,7 +96,7 @@ class AsyncWrappedMessages(AsyncMessages):
         **kwargs: Any,
     ):
         if posthog_trace_id is None:
-            posthog_trace_id = uuid.uuid4()
+            posthog_trace_id = str(uuid.uuid4())
 
         return await self._create_streaming(
             posthog_distinct_id,
@@ -116,7 +123,7 @@ class AsyncWrappedMessages(AsyncMessages):
 
         async def generator():
             nonlocal usage_stats
-            nonlocal accumulated_content
+            nonlocal accumulated_content  # noqa: F824
             try:
                 async for event in response:
                     if hasattr(event, "usage") and event.usage:
@@ -125,6 +132,8 @@ class AsyncWrappedMessages(AsyncMessages):
                             for k in [
                                 "input_tokens",
                                 "output_tokens",
+                                "cache_read_input_tokens",
+                                "cache_creation_input_tokens",
                             ]
                         }
 
@@ -165,7 +174,7 @@ class AsyncWrappedMessages(AsyncMessages):
         output: str,
     ):
         if posthog_trace_id is None:
-            posthog_trace_id = uuid.uuid4()
+            posthog_trace_id = str(uuid.uuid4())
 
         event_properties = {
             "$ai_provider": "anthropic",
@@ -184,6 +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_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
 
@@ -60,6 +69,8 @@ class GenerationMetadata(SpanMetadata):
     """Model parameters of the run: temperature, max_tokens, etc."""
     base_url: Optional[str] = None
     """Base URL of the provider's API used in the run."""
+    tools: Optional[List[Dict[str, Any]]] = None
+    """Tools provided to the model."""
 
 
 RunMetadata = Union[SpanMetadata, GenerationMetadata]
@@ -141,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,
@@ -174,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(
@@ -214,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(
@@ -238,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,
@@ -277,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,
@@ -287,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(
@@ -362,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,
@@ -374,9 +403,13 @@ 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"):
+                generation.tools = tools
         if isinstance(metadata, dict):
             if model := metadata.get("ls_model_name"):
                 generation.model = model
@@ -406,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.
         """
@@ -414,17 +449,25 @@ 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, run_id, run, outputs, self._get_parent_run_id(trace_id, run_id, parent_run_id)
+            trace_id,
+            run_id,
+            run,
+            outputs,
+            self._get_parent_run_id(trace_id, run_id, parent_run_id),
         )
 
     def _capture_trace_or_span(
@@ -438,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,
@@ -452,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
@@ -465,7 +512,10 @@ class CallbackHandler(BaseCallbackHandler):
         )
 
     def _pop_run_and_capture_generation(
-        self, run_id: UUID, parent_run_id: Optional[UUID], response: Union[LLMResult, BaseException]
+        self,
+        run_id: UUID,
+        parent_run_id: Optional[UUID],
+        response: Union[LLMResult, BaseException],
     ):
         trace_id = self._get_trace_id(run_id)
         self._pop_parent_of_run(run_id)
@@ -473,10 +523,16 @@ 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, run_id, run, response, self._get_parent_run_id(trace_id, run_id, parent_run_id)
+            trace_id,
+            run_id,
+            run,
+            response,
+            self._get_parent_run_id(trace_id, run_id, parent_run_id),
         )
 
     def _capture_generation(
@@ -500,6 +556,12 @@ class CallbackHandler(BaseCallbackHandler):
             "$ai_latency": run.latency,
             "$ai_base_url": run.base_url,
         }
+        if run.tools:
+            event_properties["$ai_tools"] = with_privacy_mode(
+                self._client,
+                self._privacy_mode,
+                run.tools,
+            )
 
         if isinstance(output, BaseException):
             event_properties["$ai_http_status"] = _get_http_status(output)
@@ -519,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)
@@ -610,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
@@ -635,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", {})
@@ -648,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
@@ -670,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/__init__.py

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

posthog/ai/openai/openai.py

@@ -1,14 +1,19 @@
 import time
 import uuid
-from typing import Any, Dict, Optional
+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,21 +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)
+
+        # 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 WrappedChat(openai.resources.chat.Chat):
-    _client: OpenAI
+class WrappedResponses:
+    """Wrapper for OpenAI responses that tracks usage in PostHog."""
 
-    @property
-    def completions(self):
-        return WrappedCompletions(self._client)
+    def __init__(self, client: OpenAI, original_responses):
+        self._client = client
+        self._original = original_responses
 
-
-class WrappedCompletions(openai.resources.chat.completions.Completions):
-    _client: OpenAI
+    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,
@@ -55,7 +76,7 @@ class WrappedCompletions(openai.resources.chat.completions.Completions):
         **kwargs: Any,
     ):
         if posthog_trace_id is None:
-            posthog_trace_id = uuid.uuid4()
+            posthog_trace_id = str(uuid.uuid4())
 
         if kwargs.get("stream", False):
             return self._create_streaming(
@@ -76,7 +97,7 @@ class WrappedCompletions(openai.resources.chat.completions.Completions):
             posthog_privacy_mode,
             posthog_groups,
             self._client.base_url,
-            super().create,
+            self._original.create,
             **kwargs,
         )
 
@@ -91,38 +112,51 @@ class WrappedCompletions(openai.resources.chat.completions.Completions):
     ):
         start_time = time.time()
         usage_stats: Dict[str, int] = {}
-        accumulated_content = []
-        if "stream_options" not in kwargs:
-            kwargs["stream_options"] = {}
-        kwargs["stream_options"]["include_usage"] = True
-        response = super().create(**kwargs)
+        final_content = []
+        response = self._original.create(**kwargs)
 
         def generator():
             nonlocal usage_stats
-            nonlocal accumulated_content
+            nonlocal final_content  # noqa: F824
+
             try:
                 for chunk in response:
+                    if hasattr(chunk, "type") and chunk.type == "response.completed":
+                        res = chunk.response
+                        if res.output and len(res.output) > 0:
+                            final_content.append(res.output[0])
+
                     if hasattr(chunk, "usage") and chunk.usage:
                         usage_stats = {
                             k: getattr(chunk.usage, k, 0)
                             for k in [
-                                "prompt_tokens",
-                                "completion_tokens",
+                                "input_tokens",
+                                "output_tokens",
                                 "total_tokens",
                             ]
                         }
 
-                    if hasattr(chunk, "choices") and chunk.choices and len(chunk.choices) > 0:
-                        content = chunk.choices[0].delta.content
-                        if content:
-                            accumulated_content.append(content)
+                        # Add support for cached tokens
+                        if hasattr(chunk.usage, "output_tokens_details") and hasattr(
+                            chunk.usage.output_tokens_details, "reasoning_tokens"
+                        ):
+                            usage_stats["reasoning_tokens"] = (
+                                chunk.usage.output_tokens_details.reasoning_tokens
+                            )
+
+                        if hasattr(chunk.usage, "input_tokens_details") and hasattr(
+                            chunk.usage.input_tokens_details, "cached_tokens"
+                        ):
+                            usage_stats["cache_read_input_tokens"] = (
+                                chunk.usage.input_tokens_details.cached_tokens
+                            )
 
                     yield chunk
 
             finally:
                 end_time = time.time()
                 latency = end_time - start_time
-                output = "".join(accumulated_content)
+                output = final_content
                 self._capture_streaming_event(
                     posthog_distinct_id,
                     posthog_trace_id,
@@ -147,30 +181,44 @@ class WrappedCompletions(openai.resources.chat.completions.Completions):
         kwargs: Dict[str, Any],
         usage_stats: Dict[str, int],
         latency: float,
-        output: str,
+        output: Any,
+        tool_calls: Optional[List[Dict[str, Any]]] = None,
     ):
         if posthog_trace_id is None:
-            posthog_trace_id = uuid.uuid4()
+            posthog_trace_id = str(uuid.uuid4())
 
         event_properties = {
             "$ai_provider": "openai",
             "$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("input")
+            ),
             "$ai_output_choices": with_privacy_mode(
                 self._client._ph_client,
                 posthog_privacy_mode,
-                [{"content": output, "role": "assistant"}],
+                output,
             ),
             "$ai_http_status": 200,
-            "$ai_input_tokens": usage_stats.get("prompt_tokens", 0),
-            "$ai_output_tokens": usage_stats.get("completion_tokens", 0),
+            "$ai_input_tokens": usage_stats.get("input_tokens", 0),
+            "$ai_output_tokens": usage_stats.get("output_tokens", 0),
+            "$ai_cache_read_input_tokens": usage_stats.get(
+                "cache_read_input_tokens", 0
+            ),
+            "$ai_reasoning_tokens": usage_stats.get("reasoning_tokens", 0),
             "$ai_latency": latency,
             "$ai_trace_id": posthog_trace_id,
             "$ai_base_url": str(self._client.base_url),
-            **posthog_properties,
+            **(posthog_properties or {}),
         }
 
+        if tool_calls:
+            event_properties["$ai_tools"] = with_privacy_mode(
+                self._client._ph_client,
+                posthog_privacy_mode,
+                tool_calls,
+            )
+
         if posthog_distinct_id is None:
             event_properties["$process_person_profile"] = False
 
@@ -183,8 +231,238 @@ class WrappedCompletions(openai.resources.chat.completions.Completions):
             )
 
 
-class WrappedEmbeddings(openai.resources.embeddings.Embeddings):
-    _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, self._original.completions)
+
+
+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,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: bool = False,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        if posthog_trace_id is None:
+            posthog_trace_id = str(uuid.uuid4())
+
+        if kwargs.get("stream", False):
+            return self._create_streaming(
+                posthog_distinct_id,
+                posthog_trace_id,
+                posthog_properties,
+                posthog_privacy_mode,
+                posthog_groups,
+                **kwargs,
+            )
+
+        return call_llm_and_track_usage(
+            posthog_distinct_id,
+            self._client._ph_client,
+            "openai",
+            posthog_trace_id,
+            posthog_properties,
+            posthog_privacy_mode,
+            posthog_groups,
+            self._client.base_url,
+            self._original.create,
+            **kwargs,
+        )
+
+    def _create_streaming(
+        self,
+        posthog_distinct_id: Optional[str],
+        posthog_trace_id: Optional[str],
+        posthog_properties: Optional[Dict[str, Any]],
+        posthog_privacy_mode: bool,
+        posthog_groups: Optional[Dict[str, Any]],
+        **kwargs: Any,
+    ):
+        start_time = time.time()
+        usage_stats: Dict[str, int] = {}
+        accumulated_content = []
+        accumulated_tools = {}
+        if "stream_options" not in kwargs:
+            kwargs["stream_options"] = {}
+        kwargs["stream_options"]["include_usage"] = True
+        response = self._original.create(**kwargs)
+
+        def generator():
+            nonlocal usage_stats
+            nonlocal accumulated_content  # noqa: F824
+            nonlocal accumulated_tools  # noqa: F824
+
+            try:
+                for chunk in response:
+                    if hasattr(chunk, "usage") and chunk.usage:
+                        usage_stats = {
+                            k: getattr(chunk.usage, k, 0)
+                            for k in [
+                                "prompt_tokens",
+                                "completion_tokens",
+                                "total_tokens",
+                            ]
+                        }
+
+                        # Add support for cached tokens
+                        if hasattr(chunk.usage, "prompt_tokens_details") and hasattr(
+                            chunk.usage.prompt_tokens_details, "cached_tokens"
+                        ):
+                            usage_stats["cache_read_input_tokens"] = (
+                                chunk.usage.prompt_tokens_details.cached_tokens
+                            )
+
+                        if hasattr(chunk.usage, "output_tokens_details") and hasattr(
+                            chunk.usage.output_tokens_details, "reasoning_tokens"
+                        ):
+                            usage_stats["reasoning_tokens"] = (
+                                chunk.usage.output_tokens_details.reasoning_tokens
+                            )
+
+                    if (
+                        hasattr(chunk, "choices")
+                        and chunk.choices
+                        and len(chunk.choices) > 0
+                    ):
+                        if chunk.choices[0].delta and chunk.choices[0].delta.content:
+                            content = chunk.choices[0].delta.content
+                            if content:
+                                accumulated_content.append(content)
+
+                        # Process tool calls
+                        tool_calls = getattr(chunk.choices[0].delta, "tool_calls", None)
+                        if tool_calls:
+                            for tool_call in tool_calls:
+                                index = tool_call.index
+                                if index not in accumulated_tools:
+                                    accumulated_tools[index] = tool_call
+                                else:
+                                    # Append arguments for existing tool calls
+                                    if hasattr(tool_call, "function") and hasattr(
+                                        tool_call.function, "arguments"
+                                    ):
+                                        accumulated_tools[
+                                            index
+                                        ].function.arguments += (
+                                            tool_call.function.arguments
+                                        )
+
+                    yield chunk
+
+            finally:
+                end_time = time.time()
+                latency = end_time - start_time
+                output = "".join(accumulated_content)
+                tools = list(accumulated_tools.values()) if accumulated_tools else None
+                self._capture_streaming_event(
+                    posthog_distinct_id,
+                    posthog_trace_id,
+                    posthog_properties,
+                    posthog_privacy_mode,
+                    posthog_groups,
+                    kwargs,
+                    usage_stats,
+                    latency,
+                    output,
+                    tools,
+                )
+
+        return generator()
+
+    def _capture_streaming_event(
+        self,
+        posthog_distinct_id: Optional[str],
+        posthog_trace_id: Optional[str],
+        posthog_properties: Optional[Dict[str, Any]],
+        posthog_privacy_mode: bool,
+        posthog_groups: Optional[Dict[str, Any]],
+        kwargs: Dict[str, Any],
+        usage_stats: Dict[str, int],
+        latency: float,
+        output: Any,
+        tool_calls: Optional[List[Dict[str, Any]]] = None,
+    ):
+        if posthog_trace_id is None:
+            posthog_trace_id = str(uuid.uuid4())
+
+        event_properties = {
+            "$ai_provider": "openai",
+            "$ai_model": kwargs.get("model"),
+            "$ai_model_parameters": get_model_params(kwargs),
+            "$ai_input": with_privacy_mode(
+                self._client._ph_client, posthog_privacy_mode, kwargs.get("messages")
+            ),
+            "$ai_output_choices": with_privacy_mode(
+                self._client._ph_client,
+                posthog_privacy_mode,
+                [{"content": output, "role": "assistant"}],
+            ),
+            "$ai_http_status": 200,
+            "$ai_input_tokens": usage_stats.get("prompt_tokens", 0),
+            "$ai_output_tokens": usage_stats.get("completion_tokens", 0),
+            "$ai_cache_read_input_tokens": usage_stats.get(
+                "cache_read_input_tokens", 0
+            ),
+            "$ai_reasoning_tokens": usage_stats.get("reasoning_tokens", 0),
+            "$ai_latency": latency,
+            "$ai_trace_id": posthog_trace_id,
+            "$ai_base_url": str(self._client.base_url),
+            **(posthog_properties or {}),
+        }
+
+        if tool_calls:
+            event_properties["$ai_tools"] = with_privacy_mode(
+                self._client._ph_client,
+                posthog_privacy_mode,
+                tool_calls,
+            )
+
+        if posthog_distinct_id is None:
+            event_properties["$process_person_profile"] = False
+
+        if hasattr(self._client._ph_client, "capture"):
+            self._client._ph_client.capture(
+                distinct_id=posthog_distinct_id or posthog_trace_id,
+                event="$ai_generation",
+                properties=event_properties,
+                groups=posthog_groups,
+            )
+
+
+class WrappedEmbeddings:
+    """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,
@@ -202,16 +480,18 @@ 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:
             The response from OpenAI's embeddings.create call.
         """
         if posthog_trace_id is None:
-            posthog_trace_id = uuid.uuid4()
+            posthog_trace_id = str(uuid.uuid4())
 
         start_time = time.time()
-        response = super().create(**kwargs)
+        response = self._original.create(**kwargs)
         end_time = time.time()
 
         # Extract usage statistics if available
@@ -228,13 +508,15 @@ 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,
             "$ai_trace_id": posthog_trace_id,
             "$ai_base_url": str(self._client.base_url),
-            **posthog_properties,
+            **(posthog_properties or {}),
         }
 
         if posthog_distinct_id is None:
@@ -252,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,
@@ -289,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

@@ -1,14 +1,19 @@
 import time
 import uuid
-from typing import Any, Dict, Optional
+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,26 +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)
+
+        # 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 WrappedChat(openai.resources.chat.AsyncChat):
-    _client: AsyncOpenAI
+class WrappedResponses:
+    """Async wrapper for OpenAI responses that tracks usage in PostHog."""
 
-    @property
-    def completions(self):
-        return WrappedCompletions(self._client)
+    def __init__(self, client: AsyncOpenAI, original_responses):
+        self._client = client
+        self._original = original_responses
 
-
-class WrappedCompletions(openai.resources.chat.completions.AsyncCompletions):
-    _client: AsyncOpenAI
+    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,
@@ -54,9 +76,8 @@ class WrappedCompletions(openai.resources.chat.completions.AsyncCompletions):
         **kwargs: Any,
     ):
         if posthog_trace_id is None:
-            posthog_trace_id = uuid.uuid4()
+            posthog_trace_id = str(uuid.uuid4())
 
-        # If streaming, handle streaming specifically
         if kwargs.get("stream", False):
             return await self._create_streaming(
                 posthog_distinct_id,
@@ -67,59 +88,75 @@ class WrappedCompletions(openai.resources.chat.completions.AsyncCompletions):
                 **kwargs,
             )
 
-        response = await call_llm_and_track_usage_async(
+        return await call_llm_and_track_usage_async(
             posthog_distinct_id,
             self._client._ph_client,
             "openai",
             posthog_trace_id,
             posthog_properties,
+            posthog_privacy_mode,
+            posthog_groups,
             self._client.base_url,
-            super().create,
+            self._original.create,
             **kwargs,
         )
-        return response
 
     async def _create_streaming(
         self,
         posthog_distinct_id: Optional[str],
         posthog_trace_id: Optional[str],
         posthog_properties: Optional[Dict[str, Any]],
-        posthog_privacy_mode: bool = False,
-        posthog_groups: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: bool,
+        posthog_groups: Optional[Dict[str, Any]],
         **kwargs: Any,
     ):
         start_time = time.time()
         usage_stats: Dict[str, int] = {}
-        accumulated_content = []
-        if "stream_options" not in kwargs:
-            kwargs["stream_options"] = {}
-        kwargs["stream_options"]["include_usage"] = True
-        response = await super().create(**kwargs)
+        final_content = []
+        response = await self._original.create(**kwargs)
 
         async def async_generator():
-            nonlocal usage_stats, accumulated_content
+            nonlocal usage_stats
+            nonlocal final_content  # noqa: F824
+
             try:
                 async for chunk in response:
+                    if hasattr(chunk, "type") and chunk.type == "response.completed":
+                        res = chunk.response
+                        if res.output and len(res.output) > 0:
+                            final_content.append(res.output[0])
+
                     if hasattr(chunk, "usage") and chunk.usage:
                         usage_stats = {
                             k: getattr(chunk.usage, k, 0)
                             for k in [
-                                "prompt_tokens",
-                                "completion_tokens",
+                                "input_tokens",
+                                "output_tokens",
                                 "total_tokens",
                             ]
                         }
-                    if hasattr(chunk, "choices") and chunk.choices and len(chunk.choices) > 0:
-                        content = chunk.choices[0].delta.content
-                        if content:
-                            accumulated_content.append(content)
+
+                        # Add support for cached tokens
+                        if hasattr(chunk.usage, "output_tokens_details") and hasattr(
+                            chunk.usage.output_tokens_details, "reasoning_tokens"
+                        ):
+                            usage_stats["reasoning_tokens"] = (
+                                chunk.usage.output_tokens_details.reasoning_tokens
+                            )
+
+                        if hasattr(chunk.usage, "input_tokens_details") and hasattr(
+                            chunk.usage.input_tokens_details, "cached_tokens"
+                        ):
+                            usage_stats["cache_read_input_tokens"] = (
+                                chunk.usage.input_tokens_details.cached_tokens
+                            )
 
                     yield chunk
 
             finally:
                 end_time = time.time()
                 latency = end_time - start_time
-                output = "".join(accumulated_content)
+                output = final_content
                 await self._capture_streaming_event(
                     posthog_distinct_id,
                     posthog_trace_id,
@@ -144,30 +181,44 @@ class WrappedCompletions(openai.resources.chat.completions.AsyncCompletions):
         kwargs: Dict[str, Any],
         usage_stats: Dict[str, int],
         latency: float,
-        output: str,
+        output: Any,
+        tool_calls: Optional[List[Dict[str, Any]]] = None,
     ):
         if posthog_trace_id is None:
-            posthog_trace_id = uuid.uuid4()
+            posthog_trace_id = str(uuid.uuid4())
 
         event_properties = {
             "$ai_provider": "openai",
             "$ai_model": kwargs.get("model"),
             "$ai_model_parameters": get_model_params(kwargs),
-            "$ai_input": with_privacy_mode(self._client._ph_client, posthog_privacy_mode, kwargs.get("messages")),
+            "$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,
-                [{"content": output, "role": "assistant"}],
+                output,
             ),
             "$ai_http_status": 200,
-            "$ai_input_tokens": usage_stats.get("prompt_tokens", 0),
-            "$ai_output_tokens": usage_stats.get("completion_tokens", 0),
+            "$ai_input_tokens": usage_stats.get("input_tokens", 0),
+            "$ai_output_tokens": usage_stats.get("output_tokens", 0),
+            "$ai_cache_read_input_tokens": usage_stats.get(
+                "cache_read_input_tokens", 0
+            ),
+            "$ai_reasoning_tokens": usage_stats.get("reasoning_tokens", 0),
             "$ai_latency": latency,
             "$ai_trace_id": posthog_trace_id,
             "$ai_base_url": str(self._client.base_url),
-            **posthog_properties,
+            **(posthog_properties or {}),
         }
 
+        if tool_calls:
+            event_properties["$ai_tools"] = with_privacy_mode(
+                self._client._ph_client,
+                posthog_privacy_mode,
+                tool_calls,
+            )
+
         if posthog_distinct_id is None:
             event_properties["$process_person_profile"] = False
 
@@ -180,8 +231,241 @@ class WrappedCompletions(openai.resources.chat.completions.AsyncCompletions):
             )
 
 
-class WrappedEmbeddings(openai.resources.embeddings.AsyncEmbeddings):
-    _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, self._original.completions)
+
+
+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,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: bool = False,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        if posthog_trace_id is None:
+            posthog_trace_id = str(uuid.uuid4())
+
+        # If streaming, handle streaming specifically
+        if kwargs.get("stream", False):
+            return await self._create_streaming(
+                posthog_distinct_id,
+                posthog_trace_id,
+                posthog_properties,
+                posthog_privacy_mode,
+                posthog_groups,
+                **kwargs,
+            )
+
+        response = await call_llm_and_track_usage_async(
+            posthog_distinct_id,
+            self._client._ph_client,
+            "openai",
+            posthog_trace_id,
+            posthog_properties,
+            posthog_privacy_mode,
+            posthog_groups,
+            self._client.base_url,
+            self._original.create,
+            **kwargs,
+        )
+        return response
+
+    async def _create_streaming(
+        self,
+        posthog_distinct_id: Optional[str],
+        posthog_trace_id: Optional[str],
+        posthog_properties: Optional[Dict[str, Any]],
+        posthog_privacy_mode: bool,
+        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 self._original.create(**kwargs)
+
+        async def async_generator():
+            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:
+                        usage_stats = {
+                            k: getattr(chunk.usage, k, 0)
+                            for k in [
+                                "prompt_tokens",
+                                "completion_tokens",
+                                "total_tokens",
+                            ]
+                        }
+
+                        # Add support for cached tokens
+                        if hasattr(chunk.usage, "prompt_tokens_details") and hasattr(
+                            chunk.usage.prompt_tokens_details, "cached_tokens"
+                        ):
+                            usage_stats["cache_read_input_tokens"] = (
+                                chunk.usage.prompt_tokens_details.cached_tokens
+                            )
+
+                        if hasattr(chunk.usage, "output_tokens_details") and hasattr(
+                            chunk.usage.output_tokens_details, "reasoning_tokens"
+                        ):
+                            usage_stats["reasoning_tokens"] = (
+                                chunk.usage.output_tokens_details.reasoning_tokens
+                            )
+
+                    if (
+                        hasattr(chunk, "choices")
+                        and chunk.choices
+                        and len(chunk.choices) > 0
+                    ):
+                        if chunk.choices[0].delta and chunk.choices[0].delta.content:
+                            content = chunk.choices[0].delta.content
+                            if content:
+                                accumulated_content.append(content)
+
+                        # Process tool calls
+                        tool_calls = getattr(chunk.choices[0].delta, "tool_calls", None)
+                        if tool_calls:
+                            for tool_call in tool_calls:
+                                index = tool_call.index
+                                if index not in accumulated_tools:
+                                    accumulated_tools[index] = tool_call
+                                else:
+                                    # Append arguments for existing tool calls
+                                    if hasattr(tool_call, "function") and hasattr(
+                                        tool_call.function, "arguments"
+                                    ):
+                                        accumulated_tools[
+                                            index
+                                        ].function.arguments += (
+                                            tool_call.function.arguments
+                                        )
+
+                    yield chunk
+
+            finally:
+                end_time = time.time()
+                latency = end_time - start_time
+                output = "".join(accumulated_content)
+                tools = list(accumulated_tools.values()) if accumulated_tools else None
+                await self._capture_streaming_event(
+                    posthog_distinct_id,
+                    posthog_trace_id,
+                    posthog_properties,
+                    posthog_privacy_mode,
+                    posthog_groups,
+                    kwargs,
+                    usage_stats,
+                    latency,
+                    output,
+                    tools,
+                )
+
+        return async_generator()
+
+    async def _capture_streaming_event(
+        self,
+        posthog_distinct_id: Optional[str],
+        posthog_trace_id: Optional[str],
+        posthog_properties: Optional[Dict[str, Any]],
+        posthog_privacy_mode: bool,
+        posthog_groups: Optional[Dict[str, Any]],
+        kwargs: Dict[str, Any],
+        usage_stats: Dict[str, int],
+        latency: float,
+        output: Any,
+        tool_calls: Optional[List[Dict[str, Any]]] = None,
+    ):
+        if posthog_trace_id is None:
+            posthog_trace_id = str(uuid.uuid4())
+
+        event_properties = {
+            "$ai_provider": "openai",
+            "$ai_model": kwargs.get("model"),
+            "$ai_model_parameters": get_model_params(kwargs),
+            "$ai_input": with_privacy_mode(
+                self._client._ph_client, posthog_privacy_mode, kwargs.get("messages")
+            ),
+            "$ai_output_choices": with_privacy_mode(
+                self._client._ph_client,
+                posthog_privacy_mode,
+                [{"content": output, "role": "assistant"}],
+            ),
+            "$ai_http_status": 200,
+            "$ai_input_tokens": usage_stats.get("prompt_tokens", 0),
+            "$ai_output_tokens": usage_stats.get("completion_tokens", 0),
+            "$ai_cache_read_input_tokens": usage_stats.get(
+                "cache_read_input_tokens", 0
+            ),
+            "$ai_reasoning_tokens": usage_stats.get("reasoning_tokens", 0),
+            "$ai_latency": latency,
+            "$ai_trace_id": posthog_trace_id,
+            "$ai_base_url": str(self._client.base_url),
+            **(posthog_properties or {}),
+        }
+
+        if tool_calls:
+            event_properties["$ai_tools"] = with_privacy_mode(
+                self._client._ph_client,
+                posthog_privacy_mode,
+                tool_calls,
+            )
+
+        if posthog_distinct_id is None:
+            event_properties["$process_person_profile"] = False
+
+        if hasattr(self._client._ph_client, "capture"):
+            self._client._ph_client.capture(
+                distinct_id=posthog_distinct_id or posthog_trace_id,
+                event="$ai_generation",
+                properties=event_properties,
+                groups=posthog_groups,
+            )
+
+
+class WrappedEmbeddings:
+    """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,
@@ -199,18 +483,18 @@ 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:
             The response from OpenAI's embeddings.create call.
         """
         if posthog_trace_id is None:
-            posthog_trace_id = uuid.uuid4()
+            posthog_trace_id = str(uuid.uuid4())
 
         start_time = time.time()
-        response = await super().create(**kwargs)
+        response = await self._original.create(**kwargs)
         end_time = time.time()
 
         # Extract usage statistics if available
@@ -227,13 +511,15 @@ 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,
             "$ai_trace_id": posthog_trace_id,
             "$ai_base_url": str(self._client.base_url),
-            **posthog_properties,
+            **(posthog_properties or {}),
         }
 
         if posthog_distinct_id is None:
@@ -251,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,
@@ -288,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

@@ -0,0 +1,95 @@
+try:
+    import openai
+except ImportError:
+    raise ModuleNotFoundError(
+        "Please install the Open AI SDK to use this feature: 'pip install openai'"
+    )
+
+from posthog.ai.openai.openai import (
+    WrappedBeta,
+    WrappedChat,
+    WrappedEmbeddings,
+    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
+
+
+class AzureOpenAI(openai.AzureOpenAI):
+    """
+    A wrapper around the Azure OpenAI SDK that automatically sends LLM usage events to PostHog.
+    """
+
+    _ph_client: PostHogClient
+
+    def __init__(self, posthog_client: PostHogClient, **kwargs):
+        """
+        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
+
+        # 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):
+    """
+    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
+
+        # 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

@@ -1,6 +1,6 @@
 import time
 import uuid
-from typing import Any, Callable, Dict, Optional
+from typing import Any, Callable, Dict, List, Optional
 
 from httpx import URL
 
@@ -34,15 +34,68 @@ def get_usage(response, provider: str) -> Dict[str, Any]:
         return {
             "input_tokens": response.usage.input_tokens,
             "output_tokens": response.usage.output_tokens,
+            "cache_read_input_tokens": response.usage.cache_read_input_tokens,
+            "cache_creation_input_tokens": response.usage.cache_creation_input_tokens,
         }
     elif provider == "openai":
+        cached_tokens = 0
+        input_tokens = 0
+        output_tokens = 0
+        reasoning_tokens = 0
+
+        # responses api
+        if hasattr(response.usage, "input_tokens"):
+            input_tokens = response.usage.input_tokens
+        if hasattr(response.usage, "output_tokens"):
+            output_tokens = response.usage.output_tokens
+        if hasattr(response.usage, "input_tokens_details") and hasattr(
+            response.usage.input_tokens_details, "cached_tokens"
+        ):
+            cached_tokens = response.usage.input_tokens_details.cached_tokens
+        if hasattr(response.usage, "output_tokens_details") and hasattr(
+            response.usage.output_tokens_details, "reasoning_tokens"
+        ):
+            reasoning_tokens = response.usage.output_tokens_details.reasoning_tokens
+
+        # chat completions
+        if hasattr(response.usage, "prompt_tokens"):
+            input_tokens = response.usage.prompt_tokens
+        if hasattr(response.usage, "completion_tokens"):
+            output_tokens = response.usage.completion_tokens
+        if hasattr(response.usage, "prompt_tokens_details") and hasattr(
+            response.usage.prompt_tokens_details, "cached_tokens"
+        ):
+            cached_tokens = response.usage.prompt_tokens_details.cached_tokens
+
         return {
-            "input_tokens": response.usage.prompt_tokens,
-            "output_tokens": response.usage.completion_tokens,
+            "input_tokens": input_tokens,
+            "output_tokens": output_tokens,
+            "cache_read_input_tokens": cached_tokens,
+            "reasoning_tokens": reasoning_tokens,
+        }
+    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,
+        "cache_read_input_tokens": 0,
+        "cache_creation_input_tokens": 0,
+        "reasoning_tokens": 0,
     }
 
 
@@ -57,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
 
 
@@ -75,24 +130,184 @@ def format_response_anthropic(response):
 
 def format_response_openai(response):
     output = []
-    for choice in response.choices:
-        if choice.message.content:
-            output.append(
-                {
-                    "content": choice.message.content,
-                    "role": choice.message.role,
-                }
-            )
+    if hasattr(response, "choices"):
+        for choice in response.choices:
+            # Handle Chat Completions response format
+            if hasattr(choice, "message") and choice.message and choice.message.content:
+                output.append(
+                    {
+                        "content": choice.message.content,
+                        "role": choice.message.role,
+                    }
+                )
+    # Handle Responses API format
+    if hasattr(response, "output"):
+        for item in response.output:
+            if item.type == "message":
+                # Extract text content from the content list
+                if hasattr(item, "content") and isinstance(item.content, list):
+                    for content_item in item.content:
+                        if (
+                            hasattr(content_item, "type")
+                            and content_item.type == "output_text"
+                            and hasattr(content_item, "text")
+                        ):
+                            output.append(
+                                {
+                                    "content": content_item.text,
+                                    "role": item.role,
+                                }
+                            )
+                        elif hasattr(content_item, "text"):
+                            output.append(
+                                {
+                                    "content": content_item.text,
+                                    "role": item.role,
+                                }
+                            )
+                        elif (
+                            hasattr(content_item, "type")
+                            and content_item.type == "input_image"
+                            and hasattr(content_item, "image_url")
+                        ):
+                            output.append(
+                                {
+                                    "content": {
+                                        "type": "image",
+                                        "image": content_item.image_url,
+                                    },
+                                    "role": item.role,
+                                }
+                            )
+                else:
+                    output.append(
+                        {
+                            "content": item.content,
+                            "role": item.role,
+                        }
+                    )
     return output
 
 
+def format_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:
+            return response.tools
+    elif provider == "openai":
+        # Handle both Chat Completions and Responses API
+        if hasattr(response, "choices") and response.choices:
+            # Check for tool_calls in message (Chat Completions format)
+            if (
+                hasattr(response.choices[0], "message")
+                and hasattr(response.choices[0].message, "tool_calls")
+                and response.choices[0].message.tool_calls
+            ):
+                return response.choices[0].message.tool_calls
+
+            # Check for tool_calls directly in response (Responses API format)
+            if (
+                hasattr(response.choices[0], "tool_calls")
+                and response.choices[0].tool_calls
+            ):
+                return response.choices[0].tool_calls
+    return None
+
+
 def merge_system_prompt(kwargs: Dict[str, Any], provider: str):
-    if provider != "anthropic":
-        return kwargs.get("messages")
-    messages = kwargs.get("messages") or []
-    if kwargs.get("system") is None:
-        return messages
-    return [{"role": "system", "content": kwargs.get("system")}] + messages
+    messages: List[Dict[str, Any]] = []
+    if provider == "anthropic":
+        messages = kwargs.get("messages") or []
+        if kwargs.get("system") is None:
+            return messages
+        return [{"role": "system", "content": kwargs.get("system")}] + messages
+    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:
+        messages = list(kwargs.get("messages", []))
+
+    if kwargs.get("input") is not None:
+        input_data = kwargs.get("input")
+        if isinstance(input_data, list):
+            messages.extend(input_data)
+        else:
+            messages.append({"role": "user", "content": input_data})
+
+    # Check if system prompt is provided as a separate parameter
+    if kwargs.get("system") is not None:
+        has_system = any(msg.get("role") == "system" for msg in messages)
+        if not has_system:
+            messages = [{"role": "system", "content": kwargs.get("system")}] + messages
+
+    # For Responses API, add instructions to the system prompt if provided
+    if kwargs.get("instructions") is not None:
+        # Find the system message if it exists
+        system_idx = next(
+            (i for i, msg in enumerate(messages) if msg.get("role") == "system"), None
+        )
+
+        if system_idx is not None:
+            # Append instructions to existing system message
+            system_content = messages[system_idx].get("content", "")
+            messages[system_idx]["content"] = (
+                f"{system_content}\n\n{kwargs.get('instructions')}"
+            )
+        else:
+            # Create a new system message with instructions
+            messages = [
+                {"role": "system", "content": kwargs.get("instructions")}
+            ] + messages
+
+    return messages
 
 
 def call_llm_and_track_usage(
@@ -122,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__(),
@@ -132,9 +349,12 @@ def call_llm_and_track_usage(
         latency = end_time - start_time
 
         if posthog_trace_id is None:
-            posthog_trace_id = uuid.uuid4()
+            posthog_trace_id = str(uuid.uuid4())
 
-        if response and hasattr(response, "usage"):
+        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)
@@ -157,9 +377,43 @@ def call_llm_and_track_usage(
             **(error_params or {}),
         }
 
+        tool_calls = format_tool_calls(response, provider)
+        if tool_calls:
+            event_properties["$ai_tools"] = with_privacy_mode(
+                ph_client, posthog_privacy_mode, tool_calls
+            )
+
+        if (
+            usage.get("cache_read_input_tokens") is not None
+            and usage.get("cache_read_input_tokens", 0) > 0
+        ):
+            event_properties["$ai_cache_read_input_tokens"] = usage.get(
+                "cache_read_input_tokens", 0
+            )
+
+        if (
+            usage.get("cache_creation_input_tokens") is not None
+            and usage.get("cache_creation_input_tokens", 0) > 0
+        ):
+            event_properties["$ai_cache_creation_input_tokens"] = usage.get(
+                "cache_creation_input_tokens", 0
+            )
+
+        if (
+            usage.get("reasoning_tokens") is not None
+            and usage.get("reasoning_tokens", 0) > 0
+        ):
+            event_properties["$ai_reasoning_tokens"] = usage.get("reasoning_tokens", 0)
+
         if posthog_distinct_id is None:
             event_properties["$process_person_profile"] = False
 
+        # Process instructions for Responses API
+        if provider == "openai" and kwargs.get("instructions") is not None:
+            event_properties["$ai_instructions"] = with_privacy_mode(
+                ph_client, posthog_privacy_mode, kwargs.get("instructions")
+            )
+
         # send the event to posthog
         if hasattr(ph_client, "capture") and callable(ph_client.capture):
             ph_client.capture(
@@ -198,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__(),
@@ -208,9 +464,12 @@ async def call_llm_and_track_usage_async(
         latency = end_time - start_time
 
         if posthog_trace_id is None:
-            posthog_trace_id = uuid.uuid4()
+            posthog_trace_id = str(uuid.uuid4())
 
-        if response and hasattr(response, "usage"):
+        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)
@@ -233,9 +492,37 @@ async def call_llm_and_track_usage_async(
             **(error_params or {}),
         }
 
+        tool_calls = format_tool_calls(response, provider)
+        if tool_calls:
+            event_properties["$ai_tools"] = with_privacy_mode(
+                ph_client, posthog_privacy_mode, tool_calls
+            )
+
+        if (
+            usage.get("cache_read_input_tokens") is not None
+            and usage.get("cache_read_input_tokens", 0) > 0
+        ):
+            event_properties["$ai_cache_read_input_tokens"] = usage.get(
+                "cache_read_input_tokens", 0
+            )
+
+        if (
+            usage.get("cache_creation_input_tokens") is not None
+            and usage.get("cache_creation_input_tokens", 0) > 0
+        ):
+            event_properties["$ai_cache_creation_input_tokens"] = usage.get(
+                "cache_creation_input_tokens", 0
+            )
+
         if posthog_distinct_id is None:
             event_properties["$process_person_profile"] = False
 
+        # Process instructions for Responses API
+        if provider == "openai" and kwargs.get("instructions") is not None:
+            event_properties["$ai_instructions"] = with_privacy_mode(
+                ph_client, posthog_privacy_mode, kwargs.get("instructions")
+            )
+
         # send the event to posthog
         if hasattr(ph_client, "capture") and callable(ph_client.capture):
             ph_client.capture(

posthog/consumer.py

@@ -1,9 +1,9 @@
 import json
 import logging
+import time
 from threading import Thread
 
 import backoff
-import monotonic
 
 from posthog.request import APIError, DatetimeSerializer, batch_post
 
@@ -96,18 +96,20 @@ class Consumer(Thread):
         queue = self.queue
         items = []
 
-        start_time = monotonic.monotonic()
+        start_time = time.monotonic()
         total_size = 0
 
         while len(items) < self.flush_at:
-            elapsed = monotonic.monotonic() - start_time
+            elapsed = time.monotonic() - start_time
             if elapsed >= self.flush_interval:
                 break
             try:
                 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

@@ -1,3 +1,8 @@
+# Portions of this file are derived from getsentry/sentry-javascript by Software, Inc. dba Sentry
+# Licensed under the MIT License
+
+# 💖open source (under MIT License)
+
 import logging
 import sys
 import threading
@@ -17,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

@@ -1,3 +1,8 @@
+# Portions of this file are derived from getsentry/sentry-javascript by Software, Inc. dba Sentry
+# Licensed under the MIT License
+
+# 💖open source (under MIT License)
+
 import re
 import sys
 from typing import TYPE_CHECKING
@@ -27,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.")
 
@@ -55,7 +59,6 @@ class DjangoIntegration:
 
 
 class DjangoRequestExtractor:
-
     def __init__(self, request):
         # type: (Any) -> None
         self.request = request
@@ -64,8 +67,8 @@ class DjangoRequestExtractor:
         headers = self.headers()
 
         # Extract traceparent and tracestate headers
-        traceparent = headers.get("traceparent")
-        tracestate = headers.get("tracestate")
+        traceparent = headers.get("Traceparent")
+        tracestate = headers.get("Tracestate")
 
         # Extract the distinct_id from tracestate
         distinct_id = None
@@ -77,12 +80,38 @@ class DjangoRequestExtractor:
                 distinct_id = match.group(1)
 
         return {
+            **self.user(),
             "distinct_id": distinct_id,
             "ip": headers.get("X-Forwarded-For"),
             "user_agent": headers.get("User-Agent"),
             "traceparent": traceparent,
+            "$request_path": self.request.path,
         }
 
+    def user(self):
+        user_data: dict[str, str] = {}
+
+        user = getattr(self.request, "user", None)
+
+        if user is None or not user.is_authenticated:
+            return user_data
+
+        try:
+            user_id = str(user.pk)
+            if user_id:
+                user_data.setdefault("$user_id", user_id)
+        except Exception:
+            pass
+
+        try:
+            email = str(user.email)
+            if email:
+                user_data.setdefault("email", email)
+        except Exception:
+            pass
+
+        return user_data
+
     def headers(self):
         # type: () -> Dict[str, str]
         return dict(self.request.headers)

posthog/exception_utils.py

@@ -1,3 +1,6 @@
+# Portions of this file are derived from getsentry/sentry-javascript by Software, Inc. dba Sentry
+# Licensed under the MIT License
+
 # copied and adapted from https://github.com/getsentry/sentry-python/blob/269d96d6e9821122fbff280e6a26956e5ed03c0b/sentry_sdk/utils.py#L689
 # 💖open source (under MIT License)
 # We want to keep payloads as similar to Sentry as possible for easy interoperability
@@ -21,7 +24,6 @@ DEFAULT_MAX_VALUE_LENGTH = 1024
 
 
 if TYPE_CHECKING:
-
     from types import FrameType, TracebackType
     from typing import (  # noqa: F401
         Any,
@@ -49,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,
@@ -57,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,
@@ -75,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
@@ -270,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)
@@ -302,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, []
 
 
@@ -339,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
 
@@ -428,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(
@@ -449,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
 
@@ -459,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
@@ -472,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
 
@@ -520,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
@@ -580,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(
@@ -601,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(
@@ -642,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(
@@ -658,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()
 
@@ -786,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,
     )
@@ -810,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

@@ -8,6 +8,7 @@ from dateutil import parser
 from dateutil.relativedelta import relativedelta
 
 from posthog import utils
+from posthog.types import FlagValue
 from posthog.utils import convert_to_datetime_aware, is_valid_regex
 
 __LONG_SCALE__ = float(0xFFFFFFFFFFFFFFF)
@@ -25,7 +26,7 @@ class InconclusiveMatchError(Exception):
 # Given the same distinct_id and key, it'll always return the same float. These floats are
 # uniformly distributed between 0 and 1, so if we want to show this feature to 20% of traffic
 # we can do _hash(key, distinct_id) < 0.2
-def _hash(key, distinct_id, salt=""):
+def _hash(key: str, distinct_id: str, salt: str = "") -> float:
     hash_key = f"{key}.{distinct_id}{salt}"
     hash_val = int(hashlib.sha1(hash_key.encode("utf-8")).hexdigest()[:15], 16)
     return hash_val / __LONG_SCALE__
@@ -42,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):
+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
@@ -69,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
@@ -80,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):
+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"):
@@ -102,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
@@ -116,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")
@@ -130,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":
@@ -148,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,
@@ -190,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)
@@ -217,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}")
@@ -235,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)
@@ -271,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"
 
@@ -302,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

@@ -7,11 +7,20 @@ from typing import Any, Optional, Union
 
 import requests
 from dateutil.tz import tzutc
+from urllib3.util.retry import Retry
 
 from posthog.utils import remove_trailing_slash
 from posthog.version import VERSION
 
-adapter = requests.adapters.HTTPAdapter(max_retries=2)
+# Retry on both connect and read errors
+# by default read errors will only retry idempotent HTTP methods (so not POST)
+adapter = requests.adapters.HTTPAdapter(
+    max_retries=Retry(
+        total=2,
+        connect=2,
+        read=2,
+    )
+)
 _session = requests.sessions.Session()
 _session.mount("https://", adapter)
 
@@ -34,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")
@@ -43,7 +57,7 @@ def post(
     url = remove_trailing_slash(host or DEFAULT_HOST) + path
     body["api_key"] = api_key
     data = json.dumps(body, cls=DatetimeSerializer)
-    log.debug("making request: %s", data)
+    log.debug("making request: %s to url: %s", data, url)
     headers = {"Content-Type": "application/json", "User-Agent": USER_AGENT}
     if gzip:
         headers["Content-Encoding"] = "gzip"
@@ -68,7 +82,21 @@ def _process_response(
     log = logging.getLogger("posthog")
     if res.status_code == 200:
         log.debug(success_message)
-        return res.json() if return_json else res
+        response = res.json() if return_json else res
+        # Handle quota limited decide responses by raising a specific error
+        # NB: other services also put entries into the quotaLimited key, but right now we only care about feature flags
+        # since most of the other services handle quota limiting in other places in the application.
+        if (
+            isinstance(response, dict)
+            and "quotaLimited" in response
+            and isinstance(response["quotaLimited"], list)
+            and "feature_flags" in response["quotaLimited"]
+        ):
+            log.warning(
+                "[FEATURE FLAGS] PostHog feature flags quota limited, resetting feature flag data.  Learn more about billing limits at https://posthog.com/docs/billing/limits-alerts"
+            )
+            raise QuotaLimitError(res.status_code, "Feature flags quota limited")
+        return response
     try:
         payload = res.json()
         log.debug("received response: %s", payload)
@@ -77,28 +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=3", gzip, timeout, **kwargs)
+    res = post(api_key, host, "/decide/?v=4", gzip, timeout, **kwargs)
     return _process_response(res, success_message="Feature flags decided successfully")
 
 
-def remote_config(personal_api_key: str, host: Optional[str] = None, key: str = "", timeout: int = 15) -> Any:
+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"
+    )
+
+
+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")
 
 
@@ -112,6 +179,10 @@ class APIError(Exception):
         return msg.format(self.message, self.status)
 
 
+class QuotaLimitError(APIError):
+    pass
+
+
 class DatetimeSerializer(json.JSONEncoder):
     def default(self, obj: Any):
         if isinstance(obj, (date, datetime)):

posthog/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")
+
+    """
+    from posthog import capture_exception
+
+    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:
+        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

@@ -4,7 +4,7 @@ from sentry_sdk.integrations import Integration
 from sentry_sdk.scope import add_global_event_processor
 from sentry_sdk.utils import Dsn
 
-import posthog
+from posthog import capture, host
 from posthog.request import DEFAULT_HOST
 from posthog.sentry import POSTHOG_ID_TAG
 
@@ -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"{host or DEFAULT_HOST}/person/{posthog_distinct_id}"
+                    )
 
                     properties = {
                         "$sentry_event_id": event["event_id"],
@@ -40,13 +44,14 @@ 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"] = (
                                 f"{PostHogIntegration.prefix}{PostHogIntegration.organization}/issues/?project={project_id}&query={event['event_id']}"
                             )
 
-                    posthog.capture(posthog_distinct_id, "$exception", properties)
+                    capture(posthog_distinct_id, "$exception", properties)
 
             return event

posthog/test/ai/anthropic/test_anthropic.py

@@ -3,12 +3,23 @@ import time
 from unittest.mock import patch
 
 import pytest
-from anthropic.types import Message, Usage
 
-from posthog.ai.anthropic import Anthropic, AsyncAnthropic
+try:
+    from anthropic.types import Message, Usage
+
+    from posthog.ai.anthropic import Anthropic, AsyncAnthropic
+
+    ANTHROPIC_AVAILABLE = True
+except ImportError:
+    ANTHROPIC_AVAILABLE = False
 
 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"
+)
+
 
 @pytest.fixture
 def mock_client():
@@ -55,8 +66,32 @@ def mock_anthropic_stream():
     return stream_generator()
 
 
+@pytest.fixture
+def mock_anthropic_response_with_cached_tokens():
+    # Create a mock Usage object with cached_tokens in input_tokens_details
+    usage = Usage(
+        input_tokens=20,
+        output_tokens=10,
+        cache_read_input_tokens=15,
+        cache_creation_input_tokens=2,
+    )
+
+    return Message(
+        id="msg_123",
+        type="message",
+        role="assistant",
+        content=[{"type": "text", "text": "Test response"}],
+        model="claude-3-opus-20240229",
+        usage=usage,
+        stop_reason="end_turn",
+        stop_sequence=None,
+    )
+
+
 def test_basic_completion(mock_client, mock_anthropic_response):
-    with patch("anthropic.resources.Messages.create", return_value=mock_anthropic_response):
+    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",
@@ -76,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
@@ -85,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",
@@ -122,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",
@@ -158,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",
@@ -175,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",
@@ -194,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(
@@ -268,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
@@ -278,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",
@@ -297,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")
@@ -328,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
 
@@ -339,3 +398,39 @@ def test_error(mock_client, mock_anthropic_response):
         props = call_args["properties"]
         assert props["$ai_is_error"] is True
         assert props["$ai_error"] == "Test error"
+
+
+def test_cached_tokens(mock_client, mock_anthropic_response_with_cached_tokens):
+    with patch(
+        "anthropic.resources.Messages.create",
+        return_value=mock_anthropic_response_with_cached_tokens,
+    ):
+        client = Anthropic(api_key="test-key", posthog_client=mock_client)
+        response = client.messages.create(
+            model="claude-3-opus-20240229",
+            messages=[{"role": "user", "content": "Hello"}],
+            posthog_distinct_id="test-id",
+            posthog_properties={"foo": "bar"},
+        )
+
+        assert response == mock_anthropic_response_with_cached_tokens
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+
+        assert call_args["distinct_id"] == "test-id"
+        assert call_args["event"] == "$ai_generation"
+        assert props["$ai_provider"] == "anthropic"
+        assert props["$ai_model"] == "claude-3-opus-20240229"
+        assert props["$ai_input"] == [{"role": "user", "content": "Hello"}]
+        assert props["$ai_output_choices"] == [
+            {"role": "assistant", "content": "Test response"}
+        ]
+        assert props["$ai_input_tokens"] == 20
+        assert props["$ai_output_tokens"] == 10
+        assert props["$ai_cache_read_input_tokens"] == 15
+        assert props["$ai_cache_creation_input_tokens"] == 2
+        assert props["$ai_http_status"] == 200
+        assert props["foo"] == "bar"
+        assert isinstance(props["$ai_latency"], float)

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

@@ -8,19 +8,44 @@ from typing import List, Literal, Optional, TypedDict, Union
 from unittest.mock import patch
 
 import pytest
-from langchain_anthropic.chat_models import ChatAnthropic
-from langchain_community.chat_models.fake import FakeMessagesListChatModel
-from langchain_community.llms.fake import FakeListLLM, FakeStreamingListLLM
-from langchain_core.messages import AIMessage, HumanMessage
-from langchain_core.prompts import ChatPromptTemplate
-from langchain_core.runnables import RunnableLambda
-from langchain_core.tools import tool
-from langchain_openai.chat_models import ChatOpenAI
-from langgraph.graph.state import END, START, StateGraph
-from langgraph.prebuilt import create_react_agent
 
-from posthog.ai.langchain import CallbackHandler
-from posthog.ai.langchain.callbacks import GenerationMetadata, SpanMetadata
+try:
+    from langchain_anthropic.chat_models import ChatAnthropic
+    from langchain_community.chat_models.fake import FakeMessagesListChatModel
+    from langchain_community.llms.fake import FakeListLLM, FakeStreamingListLLM
+    from langchain_core.messages import AIMessage, HumanMessage
+    from langchain_core.prompts import ChatPromptTemplate
+    from langchain_core.runnables import RunnableLambda
+    from langchain_core.tools import tool
+    from langchain_openai.chat_models import ChatOpenAI
+    from langgraph.graph.state import END, START, StateGraph
+    from langgraph.prebuilt import create_react_agent
+
+    from posthog.ai.langchain import CallbackHandler
+    from posthog.ai.langchain.callbacks import GenerationMetadata, SpanMetadata
+
+    LANGCHAIN_AVAILABLE = True
+except ImportError:
+
+    class FakeListLLM:
+        pass
+
+    class FakeStreamingListLLM:
+        pass
+
+    class HumanMessage:
+        pass
+
+    class AIMessage:
+        pass
+
+    LANGCHAIN_AVAILABLE = False
+
+
+# Skip all tests if LangChain 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")
@@ -229,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."
@@ -293,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
@@ -308,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)
@@ -329,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
@@ -344,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)
@@ -395,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
@@ -416,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)
@@ -452,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]
@@ -492,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]
@@ -551,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)
 
@@ -618,88 +677,66 @@ def test_graph_state(mock_client):
     assert isinstance(result["messages"][2], AIMessage)
     assert result["messages"][2].content == "It's a type of greeble."
 
-    assert mock_client.capture.call_count == 11
+    assert mock_client.capture.call_count == 6
     calls = [call[1] for call in mock_client.capture.call_args_list]
 
-    trace_args = calls[10]
-    trace_props = calls[10]["properties"]
+    # The trace event is captured at the end
+    trace_args = calls[-1]
+    trace_props = calls[-1]["properties"]
 
     # Events are captured in the reverse order.
     # Check all trace_ids
     for call in calls:
         assert call["properties"]["$ai_trace_id"] == trace_props["$ai_trace_id"]
 
-    # First span, write the state
-    assert calls[0]["event"] == "$ai_span"
-    assert calls[0]["properties"]["$ai_parent_id"] == calls[2]["properties"]["$ai_span_id"]
-    assert "$ai_span_id" in calls[0]["properties"]
-    assert calls[0]["properties"]["$ai_input_state"] == initial_state
-    assert calls[0]["properties"]["$ai_output_state"] == initial_state
-
-    # Second span, set the START node
-    assert calls[1]["event"] == "$ai_span"
-    assert calls[1]["properties"]["$ai_parent_id"] == calls[2]["properties"]["$ai_span_id"]
-    assert "$ai_span_id" in calls[1]["properties"]
-    assert calls[1]["properties"]["$ai_input_state"] == initial_state
-    assert calls[1]["properties"]["$ai_output_state"] == initial_state
-
-    # Third span, finish initialization
-    assert calls[2]["event"] == "$ai_span"
-    assert "$ai_span_id" in calls[2]["properties"]
-    assert calls[2]["properties"]["$ai_span_name"] == START
-    assert calls[2]["properties"]["$ai_parent_id"] == trace_props["$ai_trace_id"]
-    assert calls[2]["properties"]["$ai_input_state"] == initial_state
-    assert calls[2]["properties"]["$ai_output_state"] == initial_state
-
-    # Fourth span, save the value of fake_plain during its execution
+    # 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",
     }
+
+    # 1. Span - the fake_plain node, which doesn't do anything
+    assert calls[0]["event"] == "$ai_span"
+    assert calls[0]["properties"]["$ai_parent_id"] == trace_props["$ai_trace_id"]
+    assert "$ai_span_id" in calls[0]["properties"]
+    assert calls[0]["properties"]["$ai_span_name"] == "fake_plain"
+    assert calls[0]["properties"]["$ai_input_state"] == initial_state
+    assert calls[0]["properties"]["$ai_output_state"] == second_state
+
+    # 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 "$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 "$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_input_state"] == second_state
-    assert calls[3]["properties"]["$ai_output_state"] == second_state
+    assert calls[3]["properties"]["$ai_span_name"] == "RunnableSequence"
 
-    # Fifth span, run the fake_plain node
+    # 5. Span - the fake_llm node
     assert calls[4]["event"] == "$ai_span"
-    assert "$ai_span_id" in calls[4]["properties"]
-    assert calls[4]["properties"]["$ai_span_name"] == "fake_plain"
     assert calls[4]["properties"]["$ai_parent_id"] == trace_props["$ai_trace_id"]
-    assert calls[4]["properties"]["$ai_input_state"] == initial_state
-    assert calls[4]["properties"]["$ai_output_state"] == second_state
+    assert "$ai_span_id" in calls[4]["properties"]
+    assert calls[4]["properties"]["$ai_span_name"] == "fake_llm"
 
-    # Sixth span, chat prompt template
-    assert calls[5]["event"] == "$ai_span"
-    assert calls[5]["properties"]["$ai_parent_id"] == calls[7]["properties"]["$ai_span_id"]
-    assert "$ai_span_id" in calls[5]["properties"]
-    assert calls[5]["properties"]["$ai_span_name"] == "ChatPromptTemplate"
-
-    # 7. Generation, fake_llm
-    assert calls[6]["event"] == "$ai_generation"
-    assert calls[6]["properties"]["$ai_parent_id"] == calls[7]["properties"]["$ai_span_id"]
-    assert "$ai_span_id" in calls[6]["properties"]
-    assert calls[6]["properties"]["$ai_span_name"] == "FakeMessagesListChatModel"
-
-    # 8. Span, RunnableSequence
-    assert calls[7]["event"] == "$ai_span"
-    assert calls[7]["properties"]["$ai_parent_id"] == calls[9]["properties"]["$ai_span_id"]
-    assert "$ai_span_id" in calls[7]["properties"]
-    assert calls[7]["properties"]["$ai_span_name"] == "RunnableSequence"
-
-    # 9. Span, fake_llm write
-    assert calls[8]["event"] == "$ai_span"
-    assert calls[8]["properties"]["$ai_parent_id"] == calls[9]["properties"]["$ai_span_id"]
-    assert "$ai_span_id" in calls[8]["properties"]
-
-    # 10. Span, fake_llm node
-    assert calls[9]["event"] == "$ai_span"
-    assert calls[9]["properties"]["$ai_parent_id"] == trace_props["$ai_trace_id"]
-    assert "$ai_span_id" in calls[9]["properties"]
-    assert calls[9]["properties"]["$ai_span_name"] == "fake_llm"
-
-    # 11. Trace
+    # 6. Trace
     assert trace_args["event"] == "$ai_trace"
     assert trace_props["$ai_span_name"] == "LangGraph"
 
@@ -713,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"
 
 
@@ -743,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 == {}
 
@@ -836,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
 
@@ -1113,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
 
@@ -1168,6 +1222,61 @@ async def test_async_anthropic_streaming(mock_client):
     assert isinstance(trace_props["$ai_output_state"], AIMessage)
 
 
+def test_metadata_tools(mock_client):
+    callbacks = CallbackHandler(mock_client)
+    run_id = uuid.uuid4()
+    tools = [
+        [
+            {
+                "type": "function",
+                "function": {
+                    "name": "foo",
+                    "description": "The foo.",
+                    "parameters": {
+                        "properties": {
+                            "bar": {
+                                "description": "The bar of foo.",
+                                "type": "string",
+                            },
+                        },
+                        "required": ["query_description", "query_kind"],
+                        "type": "object",
+                        "additionalProperties": False,
+                    },
+                    "strict": True,
+                },
+            }
+        ]
+    ]
+
+    with patch("time.time", return_value=1234567890):
+        callbacks._set_llm_metadata(
+            {"kwargs": {"openai_api_base": "https://us.posthog.com"}},
+            run_id,
+            messages=[{"role": "user", "content": "What's the weather like in SF?"}],
+            invocation_params={"temperature": 0.5, "tools": tools},
+            metadata={"ls_model_name": "hog-mini", "ls_provider": "posthog"},
+            name="test",
+        )
+    expected = GenerationMetadata(
+        model="hog-mini",
+        input=[{"role": "user", "content": "What's the weather like in SF?"}],
+        start_time=1234567890,
+        model_params={"temperature": 0.5},
+        provider="posthog",
+        base_url="https://us.posthog.com",
+        name="test",
+        tools=tools,
+        end_time=None,
+    )
+    assert callbacks._runs[run_id] == expected
+    with patch("time.time", return_value=1234567891):
+        run = callbacks._pop_run_metadata(run_id)
+    expected.end_time = 1234567891
+    assert run == expected
+    assert callbacks._runs == {}
+
+
 def test_tool_calls(mock_client):
     prompt = ChatPromptTemplate.from_messages([("user", "Foo")])
     model = FakeMessagesListChatModel(
@@ -1205,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):
@@ -1227,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"
@@ -1235,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
     )
 
 
@@ -1257,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
@@ -1277,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
@@ -1286,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])
@@ -1309,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"]
 
@@ -1326,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"]
 
 
@@ -1342,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

@@ -1,14 +1,43 @@
+import json
 import time
 from unittest.mock import patch
 
 import pytest
-from openai.types.chat import ChatCompletion, ChatCompletionMessage
-from openai.types.chat.chat_completion import Choice
-from openai.types.completion_usage import CompletionUsage
-from openai.types.create_embedding_response import CreateEmbeddingResponse, Usage
-from openai.types.embedding import Embedding
 
-from posthog.ai.openai import OpenAI
+try:
+    from openai.types.chat import ChatCompletion, ChatCompletionMessage
+    from openai.types.chat.chat_completion import Choice
+    from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
+    from openai.types.chat.chat_completion_chunk import Choice as ChoiceChunk
+    from openai.types.chat.chat_completion_chunk import (
+        ChoiceDelta,
+        ChoiceDeltaToolCall,
+        ChoiceDeltaToolCallFunction,
+    )
+    from openai.types.chat.chat_completion_message_tool_call import (
+        ChatCompletionMessageToolCall,
+        Function,
+    )
+    from openai.types.completion_usage import CompletionUsage
+    from openai.types.create_embedding_response import CreateEmbeddingResponse, Usage
+    from openai.types.embedding import Embedding
+    from openai.types.responses import (
+        Response,
+        ResponseOutputMessage,
+        ResponseOutputText,
+        ResponseUsage,
+    )
+
+    from posthog.ai.openai import OpenAI
+
+    OPENAI_AVAILABLE = True
+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"
+)
 
 
 @pytest.fixture
@@ -43,6 +72,49 @@ def mock_openai_response():
     )
 
 
+@pytest.fixture
+def mock_openai_response_with_responses_api():
+    return Response(
+        id="test",
+        model="gpt-4o-mini",
+        object="response",
+        created_at=1741476542,
+        status="completed",
+        error=None,
+        incomplete_details=None,
+        instructions=None,
+        max_output_tokens=None,
+        tools=[],
+        tool_choice="auto",
+        output=[
+            ResponseOutputMessage(
+                id="msg_123",
+                type="message",
+                role="assistant",
+                status="completed",
+                content=[
+                    ResponseOutputText(
+                        type="output_text",
+                        text="Test response",
+                        annotations=[],
+                    )
+                ],
+            )
+        ],
+        parallel_tool_calls=True,
+        previous_response_id=None,
+        usage=ResponseUsage(
+            input_tokens=10,
+            output_tokens=10,
+            input_tokens_details={"prompt_tokens": 10, "cached_tokens": 0},
+            output_tokens_details={"reasoning_tokens": 15},
+            total_tokens=20,
+        ),
+        user=None,
+        metadata={},
+    )
+
+
 @pytest.fixture
 def mock_embedding_response():
     return CreateEmbeddingResponse(
@@ -62,8 +134,72 @@ def mock_embedding_response():
     )
 
 
+@pytest.fixture
+def mock_openai_response_with_cached_tokens():
+    return ChatCompletion(
+        id="test",
+        model="gpt-4",
+        object="chat.completion",
+        created=int(time.time()),
+        choices=[
+            Choice(
+                finish_reason="stop",
+                index=0,
+                message=ChatCompletionMessage(
+                    content="Test response",
+                    role="assistant",
+                ),
+            )
+        ],
+        usage=CompletionUsage(
+            completion_tokens=10,
+            prompt_tokens=20,
+            total_tokens=30,
+            prompt_tokens_details={"cached_tokens": 15},
+        ),
+    )
+
+
+@pytest.fixture
+def mock_openai_response_with_tool_calls():
+    return ChatCompletion(
+        id="test",
+        model="gpt-4",
+        object="chat.completion",
+        created=int(time.time()),
+        choices=[
+            Choice(
+                finish_reason="tool_calls",
+                index=0,
+                message=ChatCompletionMessage(
+                    content="I'll check the weather for you.",
+                    role="assistant",
+                    tool_calls=[
+                        ChatCompletionMessageToolCall(
+                            id="call_abc123",
+                            type="function",
+                            function=Function(
+                                name="get_weather",
+                                arguments='{"location": "San Francisco", "unit": "celsius"}',
+                            ),
+                        )
+                    ],
+                ),
+            )
+        ],
+        usage=CompletionUsage(
+            completion_tokens=15,
+            prompt_tokens=20,
+            total_tokens=35,
+        ),
+    )
+
+
 def test_basic_completion(mock_client, mock_openai_response):
-    with patch("openai.resources.chat.completions.Completions.create", return_value=mock_openai_response):
+    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",
@@ -83,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
@@ -92,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",
@@ -119,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",
@@ -137,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",
@@ -156,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(
@@ -176,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
 
@@ -187,3 +342,307 @@ def test_error(mock_client, mock_openai_response):
         props = call_args["properties"]
         assert props["$ai_is_error"] is True
         assert props["$ai_error"] == "Test error"
+
+
+def test_cached_tokens(mock_client, mock_openai_response_with_cached_tokens):
+    with patch(
+        "openai.resources.chat.completions.Completions.create",
+        return_value=mock_openai_response_with_cached_tokens,
+    ):
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+        response = client.chat.completions.create(
+            model="gpt-4",
+            messages=[{"role": "user", "content": "Hello"}],
+            posthog_distinct_id="test-id",
+            posthog_properties={"foo": "bar"},
+        )
+
+        assert response == mock_openai_response_with_cached_tokens
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+
+        assert call_args["distinct_id"] == "test-id"
+        assert call_args["event"] == "$ai_generation"
+        assert props["$ai_provider"] == "openai"
+        assert props["$ai_model"] == "gpt-4"
+        assert props["$ai_input"] == [{"role": "user", "content": "Hello"}]
+        assert props["$ai_output_choices"] == [
+            {"role": "assistant", "content": "Test response"}
+        ]
+        assert props["$ai_input_tokens"] == 20
+        assert props["$ai_output_tokens"] == 10
+        assert props["$ai_cache_read_input_tokens"] == 15
+        assert props["$ai_http_status"] == 200
+        assert props["foo"] == "bar"
+        assert isinstance(props["$ai_latency"], float)
+
+
+def test_tool_calls(mock_client, mock_openai_response_with_tool_calls):
+    with patch(
+        "openai.resources.chat.completions.Completions.create",
+        return_value=mock_openai_response_with_tool_calls,
+    ):
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+        response = client.chat.completions.create(
+            model="gpt-4",
+            messages=[
+                {"role": "user", "content": "What's the weather in San Francisco?"}
+            ],
+            tools=[
+                {
+                    "type": "function",
+                    "function": {
+                        "name": "get_weather",
+                        "description": "Get weather",
+                        "parameters": {},
+                    },
+                }
+            ],
+            posthog_distinct_id="test-id",
+        )
+
+        assert response == mock_openai_response_with_tool_calls
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+
+        assert call_args["distinct_id"] == "test-id"
+        assert call_args["event"] == "$ai_generation"
+        assert props["$ai_provider"] == "openai"
+        assert props["$ai_model"] == "gpt-4"
+        assert props["$ai_input"] == [
+            {"role": "user", "content": "What's the weather in San Francisco?"}
+        ]
+        assert props["$ai_output_choices"] == [
+            {"role": "assistant", "content": "I'll check the weather for you."}
+        ]
+
+        # Check that tool calls are properly captured
+        assert "$ai_tools" in props
+        tool_calls = props["$ai_tools"]
+        assert len(tool_calls) == 1
+
+        # Verify the tool call details
+        tool_call = tool_calls[0]
+        assert tool_call.id == "call_abc123"
+        assert tool_call.type == "function"
+        assert tool_call.function.name == "get_weather"
+
+        # Verify the arguments
+        arguments = tool_call.function.arguments
+        parsed_args = json.loads(arguments)
+        assert parsed_args == {"location": "San Francisco", "unit": "celsius"}
+
+        # Check token usage
+        assert props["$ai_input_tokens"] == 20
+        assert props["$ai_output_tokens"] == 15
+        assert props["$ai_http_status"] == 200
+
+
+def test_streaming_with_tool_calls(mock_client):
+    # Create mock tool call chunks that will be returned in sequence
+    tool_call_chunks = [
+        ChatCompletionChunk(
+            id="chunk1",
+            model="gpt-4",
+            object="chat.completion.chunk",
+            created=1234567890,
+            choices=[
+                ChoiceChunk(
+                    index=0,
+                    delta=ChoiceDelta(
+                        role="assistant",
+                        tool_calls=[
+                            ChoiceDeltaToolCall(
+                                index=0,
+                                id="call_abc123",
+                                type="function",
+                                function=ChoiceDeltaToolCallFunction(
+                                    name="get_weather",
+                                    arguments='{"location": "',
+                                ),
+                            )
+                        ],
+                    ),
+                    finish_reason=None,
+                )
+            ],
+        ),
+        ChatCompletionChunk(
+            id="chunk2",
+            model="gpt-4",
+            object="chat.completion.chunk",
+            created=1234567891,
+            choices=[
+                ChoiceChunk(
+                    index=0,
+                    delta=ChoiceDelta(
+                        tool_calls=[
+                            ChoiceDeltaToolCall(
+                                index=0,
+                                id="call_abc123",
+                                type="function",
+                                function=ChoiceDeltaToolCallFunction(
+                                    arguments='San Francisco"',
+                                ),
+                            )
+                        ],
+                    ),
+                    finish_reason=None,
+                )
+            ],
+        ),
+        ChatCompletionChunk(
+            id="chunk3",
+            model="gpt-4",
+            object="chat.completion.chunk",
+            created=1234567892,
+            choices=[
+                ChoiceChunk(
+                    index=0,
+                    delta=ChoiceDelta(
+                        tool_calls=[
+                            ChoiceDeltaToolCall(
+                                index=0,
+                                id="call_abc123",
+                                type="function",
+                                function=ChoiceDeltaToolCallFunction(
+                                    arguments=', "unit": "celsius"}',
+                                ),
+                            )
+                        ],
+                    ),
+                    finish_reason=None,
+                )
+            ],
+        ),
+        ChatCompletionChunk(
+            id="chunk4",
+            model="gpt-4",
+            object="chat.completion.chunk",
+            created=1234567893,
+            choices=[
+                ChoiceChunk(
+                    index=0,
+                    delta=ChoiceDelta(
+                        content="The weather in San Francisco is 15°C.",
+                    ),
+                    finish_reason=None,
+                )
+            ],
+            usage=CompletionUsage(
+                prompt_tokens=20,
+                completion_tokens=15,
+                total_tokens=35,
+            ),
+        ),
+    ]
+
+    # Mock the create method to return our chunks
+    with patch("openai.resources.chat.completions.Completions.create") as mock_create:
+        # Set up the mock to return our chunks when iterated
+        mock_create.return_value = tool_call_chunks
+
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+
+        # Call the streaming method
+        response_generator = client.chat.completions.create(
+            model="gpt-4",
+            messages=[
+                {"role": "user", "content": "What's the weather in San Francisco?"}
+            ],
+            tools=[
+                {
+                    "type": "function",
+                    "function": {
+                        "name": "get_weather",
+                        "description": "Get weather",
+                        "parameters": {},
+                    },
+                }
+            ],
+            stream=True,
+            posthog_distinct_id="test-id",
+        )
+
+        # Consume the generator to trigger the event capture
+        chunks = list(response_generator)
+
+        # Verify the chunks were returned correctly
+        assert len(chunks) == 4
+        assert chunks == tool_call_chunks
+
+        # Verify the capture was called with the right arguments
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+
+        assert call_args["distinct_id"] == "test-id"
+        assert call_args["event"] == "$ai_generation"
+        assert props["$ai_provider"] == "openai"
+        assert props["$ai_model"] == "gpt-4"
+
+        # Check that the tool calls were properly accumulated
+        assert "$ai_tools" in props
+        tool_calls = props["$ai_tools"]
+        assert len(tool_calls) == 1
+
+        # Verify the complete tool call was properly assembled
+        tool_call = tool_calls[0]
+        assert tool_call.id == "call_abc123"
+        assert tool_call.type == "function"
+        assert tool_call.function.name == "get_weather"
+
+        # Verify the arguments were concatenated correctly
+        arguments = tool_call.function.arguments
+        parsed_args = json.loads(arguments)
+        assert parsed_args == {"location": "San Francisco", "unit": "celsius"}
+
+        # Check that the content was also accumulated
+        assert (
+            props["$ai_output_choices"][0]["content"]
+            == "The weather in San Francisco is 15°C."
+        )
+
+        # Check token usage
+        assert props["$ai_input_tokens"] == 20
+        assert props["$ai_output_tokens"] == 15
+
+
+# test responses api
+def test_responses_api(mock_client, mock_openai_response_with_responses_api):
+    with patch(
+        "openai.resources.responses.Responses.create",
+        return_value=mock_openai_response_with_responses_api,
+    ):
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+        response = client.responses.create(
+            model="gpt-4o-mini",
+            input="Hello",
+            posthog_distinct_id="test-id",
+            posthog_properties={"foo": "bar"},
+        )
+        assert response == mock_openai_response_with_responses_api
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+
+        assert call_args["distinct_id"] == "test-id"
+        assert call_args["event"] == "$ai_generation"
+        assert props["$ai_provider"] == "openai"
+        assert props["$ai_model"] == "gpt-4o-mini"
+        assert props["$ai_input"] == [{"role": "user", "content": "Hello"}]
+        assert props["$ai_output_choices"] == [
+            {"role": "assistant", "content": "Test response"}
+        ]
+        assert props["$ai_input_tokens"] == 10
+        assert props["$ai_output_tokens"] == 10
+        assert props["$ai_reasoning_tokens"] == 15
+        assert props["$ai_http_status"] == 200
+        assert props["foo"] == "bar"
+        assert isinstance(props["$ai_latency"], float)

posthog/test/exception_integrations/test_django.py

@@ -1,22 +1,44 @@
 from posthog.exception_integrations.django import DjangoRequestExtractor
+from django.test import RequestFactory
+from django.conf import settings
+from django.core.management import call_command
+import django
 
-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"
+
+# setup a test app
+if not settings.configured:
+    settings.configure(
+        SECRET_KEY="test",
+        DEFAULT_CHARSET="utf-8",
+        INSTALLED_APPS=[
+            "django.contrib.auth",
+            "django.contrib.contenttypes",
+        ],
+        DATABASES={
+            "default": {
+                "ENGINE": "django.db.backends.sqlite3",
+                "NAME": ":memory:",
+            }
+        },
+    )
+    django.setup()
+
+    call_command("migrate", verbosity=0, interactive=False)
 
 
 def mock_request_factory(override_headers):
-    class Request:
-        META = {}
-        # TRICKY: Actual django request dict object has case insensitive matching, and strips http from the names
-        headers = {
+    factory = RequestFactory(
+        headers={
             "User-Agent": DEFAULT_USER_AGENT,
             "Referrer": "http://example.com",
             "X-Forwarded-For": "193.4.5.12",
             **(override_headers or {}),
         }
+    )
 
-    return Request()
+    request = factory.get("/api/endpoint")
+    return request
 
 
 def test_request_extractor_with_no_trace():
@@ -27,17 +49,22 @@ def test_request_extractor_with_no_trace():
         "user_agent": DEFAULT_USER_AGENT,
         "traceparent": None,
         "distinct_id": None,
+        "$request_path": "/api/endpoint",
     }
 
 
 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",
         "user_agent": DEFAULT_USER_AGENT,
         "traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01",
         "distinct_id": None,
+        "$request_path": "/api/endpoint",
     }
 
 
@@ -54,15 +81,41 @@ def test_request_extractor_with_tracestate():
         "user_agent": DEFAULT_USER_AGENT,
         "traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01",
         "distinct_id": "1234",
+        "$request_path": "/api/endpoint",
     }
 
 
 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",
         "user_agent": DEFAULT_USER_AGENT,
         "traceparent": None,
         "distinct_id": "alohaMountainsXUYZ",
+        "$request_path": "/api/endpoint",
+    }
+
+
+def test_request_extractor_with_request_user():
+    from django.contrib.auth.models import User
+
+    user = User.objects.create_user(
+        username="test", email="test@posthog.com", password="top_secret"
+    )
+
+    request = mock_request_factory(None)
+    request.user = user
+
+    extractor = DjangoRequestExtractor(request)
+    assert extractor.extract_person_data() == {
+        "ip": "193.4.5.12",
+        "user_agent": DEFAULT_USER_AGENT,
+        "traceparent": None,
+        "distinct_id": None,
+        "$request_path": "/api/endpoint",
+        "email": "test@posthog.com",
+        "$user_id": "1",
     }

posthog/test/test_before_send.py

@@ -0,0 +1,171 @@
+import unittest
+
+import mock
+
+from posthog.client import Client
+from posthog.test.test_utils import FAKE_TEST_API_KEY
+
+
+class TestClient(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        # This ensures no real HTTP POST requests are made
+        cls.client_post_patcher = mock.patch("posthog.client.batch_post")
+        cls.consumer_post_patcher = mock.patch("posthog.consumer.batch_post")
+        cls.client_post_patcher.start()
+        cls.consumer_post_patcher.start()
+
+    @classmethod
+    def tearDownClass(cls):
+        cls.client_post_patcher.stop()
+        cls.consumer_post_patcher.stop()
+
+    def set_fail(self, e, batch):
+        """Mark the failure handler"""
+        print("FAIL", e, batch)  # noqa: T201
+        self.failed = True
+
+    def setUp(self):
+        self.failed = False
+        self.client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail)
+
+    def test_before_send_callback_modifies_event(self):
+        """Test that before_send callback can modify events."""
+        processed_events = []
+
+        def my_before_send(event):
+            processed_events.append(event.copy())
+            if "properties" not in event:
+                event["properties"] = {}
+            event["properties"]["processed_by_before_send"] = True
+            return event
+
+        client = Client(
+            FAKE_TEST_API_KEY, on_error=self.set_fail, before_send=my_before_send
+        )
+        success, msg = client.capture("user1", "test_event", {"original": "value"})
+
+        self.assertTrue(success)
+        self.assertEqual(msg["properties"]["processed_by_before_send"], True)
+        self.assertEqual(msg["properties"]["original"], "value")
+        self.assertEqual(len(processed_events), 1)
+        self.assertEqual(processed_events[0]["event"], "test_event")
+
+    def test_before_send_callback_drops_event(self):
+        """Test that before_send callback can drop events by returning None."""
+
+        def drop_test_events(event):
+            if event.get("event") == "test_drop_me":
+                return None
+            return event
+
+        client = Client(
+            FAKE_TEST_API_KEY, on_error=self.set_fail, before_send=drop_test_events
+        )
+
+        # Event should be dropped
+        success, msg = client.capture("user1", "test_drop_me")
+        self.assertTrue(success)
+        self.assertIsNone(msg)
+
+        # Event should go through
+        success, msg = client.capture("user1", "keep_me")
+        self.assertTrue(success)
+        self.assertIsNotNone(msg)
+        self.assertEqual(msg["event"], "keep_me")
+
+    def test_before_send_callback_handles_exceptions(self):
+        """Test that exceptions in before_send don't crash the client."""
+
+        def buggy_before_send(event):
+            raise ValueError("Oops!")
+
+        client = Client(
+            FAKE_TEST_API_KEY, on_error=self.set_fail, before_send=buggy_before_send
+        )
+        success, msg = client.capture("user1", "robust_event")
+
+        # Event should still be sent despite the exception
+        self.assertTrue(success)
+        self.assertIsNotNone(msg)
+        self.assertEqual(msg["event"], "robust_event")
+
+    def test_before_send_callback_works_with_all_event_types(self):
+        """Test that before_send works with capture, identify, set, etc."""
+
+        def add_marker(event):
+            if "properties" not in event:
+                event["properties"] = {}
+            event["properties"]["marked"] = True
+            return event
+
+        client = Client(
+            FAKE_TEST_API_KEY, on_error=self.set_fail, before_send=add_marker
+        )
+
+        # Test capture
+        success, msg = client.capture("user1", "event")
+        self.assertTrue(success)
+        self.assertTrue(msg["properties"]["marked"])
+
+        # Test identify
+        success, msg = client.identify("user1", {"trait": "value"})
+        self.assertTrue(success)
+        self.assertTrue(msg["properties"]["marked"])
+
+        # Test set
+        success, msg = client.set("user1", {"prop": "value"})
+        self.assertTrue(success)
+        self.assertTrue(msg["properties"]["marked"])
+
+        # Test page
+        success, msg = client.page("user1", "https://example.com")
+        self.assertTrue(success)
+        self.assertTrue(msg["properties"]["marked"])
+
+    def test_before_send_callback_disabled_when_none(self):
+        """Test that client works normally when before_send is None."""
+        client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, before_send=None)
+        success, msg = client.capture("user1", "normal_event")
+
+        self.assertTrue(success)
+        self.assertIsNotNone(msg)
+        self.assertEqual(msg["event"], "normal_event")
+
+    def test_before_send_callback_pii_scrubbing_example(self):
+        """Test a realistic PII scrubbing use case."""
+
+        def scrub_pii(event):
+            properties = event.get("properties", {})
+
+            # Mask email but keep domain
+            if "email" in properties:
+                email = properties["email"]
+                if "@" in email:
+                    domain = email.split("@")[1]
+                    properties["email"] = f"***@{domain}"
+                else:
+                    properties["email"] = "***"
+
+            # Remove credit card
+            properties.pop("credit_card", None)
+
+            return event
+
+        client = Client(
+            FAKE_TEST_API_KEY, on_error=self.set_fail, before_send=scrub_pii
+        )
+        success, msg = client.capture(
+            "user1",
+            "form_submit",
+            {
+                "email": "user@example.com",
+                "credit_card": "1234-5678-9012-3456",
+                "form_name": "contact",
+            },
+        )
+
+        self.assertTrue(success)
+        self.assertEqual(msg["properties"]["email"], "***@example.com")
+        self.assertNotIn("credit_card", msg["properties"])
+        self.assertEqual(msg["properties"]["form_name"], "contact")

posthog/test/test_client.py

@@ -5,9 +5,12 @@ from uuid import uuid4
 
 import mock
 import six
+from parameterized import parameterized
 
 from posthog.client import Client
+from posthog.request import APIError
 from posthog.test.test_utils import FAKE_TEST_API_KEY
+from posthog.types import FeatureFlag, LegacyFlagMetadata
 from posthog.version import VERSION
 
 
@@ -53,6 +56,11 @@ class TestClient(unittest.TestCase):
         self.assertEqual(msg["distinct_id"], "distinct_id")
         self.assertEqual(msg["properties"]["$lib"], "posthog-python")
         self.assertEqual(msg["properties"]["$lib_version"], VERSION)
+        # these will change between platforms so just asssert on presence here
+        assert msg["properties"]["$python_runtime"] == mock.ANY
+        assert msg["properties"]["$python_version"] == mock.ANY
+        assert msg["properties"]["$os"] == mock.ANY
+        assert msg["properties"]["$os_version"] == mock.ANY
 
     def test_basic_capture_with_uuid(self):
         client = self.client
@@ -100,7 +108,6 @@ class TestClient(unittest.TestCase):
         self.assertEqual(msg["properties"]["source"], "repo-name")
 
     def test_basic_capture_exception(self):
-
         with mock.patch.object(Client, "capture", return_value=None) as patch_capture:
             client = self.client
             exception = Exception("test exception")
@@ -128,7 +135,6 @@ class TestClient(unittest.TestCase):
             )
 
     def test_basic_capture_exception_with_distinct_id(self):
-
         with mock.patch.object(Client, "capture", return_value=None) as patch_capture:
             client = self.client
             exception = Exception("test exception")
@@ -156,9 +162,10 @@ 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")
 
@@ -183,10 +190,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")
 
@@ -212,7 +224,6 @@ class TestClient(unittest.TestCase):
             )
 
     def test_basic_capture_exception_with_no_exception_given(self):
-
         with mock.patch.object(Client, "capture", return_value=None) as patch_capture:
             client = self.client
             try:
@@ -226,33 +237,49 @@ 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:
             with self.assertLogs("posthog", level="WARNING") as logs:
-
                 client = self.client
                 client.capture_exception()
 
@@ -262,12 +289,29 @@ class TestClient(unittest.TestCase):
                     "WARNING:posthog:No exception information available",
                 )
 
-    @mock.patch("posthog.client.decide")
-    def test_basic_capture_with_feature_flags(self, patch_decide):
-        patch_decide.return_value = {"featureFlags": {"beta-feature": "random-variant"}}
+    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"
+            )
+            self.assertEqual(getattr(logs.records[0], "path"), "one/two/three")
 
-        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)
+    @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.flush()
         self.assertTrue(success)
         self.assertFalse(self.failed)
@@ -281,12 +325,16 @@ class TestClient(unittest.TestCase):
         self.assertEqual(msg["properties"]["$feature/beta-feature"], "random-variant")
         self.assertEqual(msg["properties"]["$active_feature_flags"], ["beta-feature"])
 
-        self.assertEqual(patch_decide.call_count, 1)
+        self.assertEqual(patch_flags.call_count, 1)
 
-    @mock.patch("posthog.client.decide")
-    def test_basic_capture_with_locally_evaluated_feature_flags(self, patch_decide):
-        patch_decide.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)
+    @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,
+        )
 
         multivariate_flag = {
             "id": 1,
@@ -298,7 +346,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,
                     },
@@ -308,12 +361,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 = {
@@ -366,12 +434,16 @@ 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_decide.call_count, 0)
+        self.assertEqual(patch_flags.call_count, 0)
 
         # test that flags are not evaluated without local evaluation
         client.feature_flags = []
@@ -384,10 +456,33 @@ class TestClient(unittest.TestCase):
         assert "$feature/false-flag" not in msg["properties"]
         assert "$active_feature_flags" not in msg["properties"]
 
-    @mock.patch("posthog.client.decide")
-    def test_dont_override_capture_with_local_flags(self, patch_decide):
-        patch_decide.return_value = {"featureFlags": {"beta-feature": "random-variant"}}
-        client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, personal_api_key=FAKE_TEST_API_KEY)
+    @mock.patch("posthog.client.get")
+    def test_load_feature_flags_quota_limited(self, patch_get):
+        mock_response = {
+            "type": "quota_limited",
+            "detail": "You have exceeded your feature flag request quota",
+            "code": "payment_required",
+        }
+        patch_get.side_effect = APIError(402, mock_response["detail"])
+
+        client = Client(FAKE_TEST_API_KEY, personal_api_key="test")
+        with self.assertLogs("posthog", level="WARNING") as logs:
+            client._load_feature_flags()
+
+            self.assertEqual(client.feature_flags, [])
+            self.assertEqual(client.feature_flags_by_key, {})
+            self.assertEqual(client.group_type_mapping, {})
+            self.assertEqual(client.cohorts, {})
+            self.assertIn("PostHog feature flags quota limited", logs.output[0])
+
+    @mock.patch("posthog.client.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,
+        )
 
         multivariate_flag = {
             "id": 1,
@@ -399,7 +494,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,
                     },
@@ -409,12 +509,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 = {
@@ -442,7 +557,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)
@@ -454,21 +571,35 @@ 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"]
 
-        self.assertEqual(patch_decide.call_count, 0)
+        self.assertEqual(patch_flags.call_count, 0)
 
-    @mock.patch("posthog.client.decide")
-    def test_basic_capture_with_feature_flags_returns_active_only(self, patch_decide):
-        patch_decide.return_value = {
-            "featureFlags": {"beta-feature": "random-variant", "alpha-feature": True, "off-feature": False}
+    @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,
+            }
         }
 
-        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)
@@ -482,10 +613,13 @@ 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_decide.call_count, 1)
-        patch_decide.assert_called_with(
+        self.assertEqual(patch_flags.call_count, 1)
+        patch_flags.assert_called_with(
             "random_key",
             "https://us.i.posthog.com",
             timeout=3,
@@ -493,13 +627,19 @@ class TestClient(unittest.TestCase):
             groups={},
             person_properties=None,
             group_properties=None,
-            disable_geoip=True,
+            geoip_disable=True,
         )
 
-    @mock.patch("posthog.client.decide")
-    def test_basic_capture_with_feature_flags_and_disable_geoip_returns_correctly(self, patch_decide):
-        patch_decide.return_value = {
-            "featureFlags": {"beta-feature": "random-variant", "alpha-feature": True, "off-feature": False}
+    @mock.patch("posthog.client.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,
+            }
         }
 
         client = Client(
@@ -510,7 +650,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)
@@ -524,10 +669,13 @@ 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_decide.call_count, 1)
-        patch_decide.assert_called_with(
+        self.assertEqual(patch_flags.call_count, 1)
+        patch_flags.assert_called_with(
             "random_key",
             "https://us.i.posthog.com",
             timeout=12,
@@ -535,15 +683,23 @@ class TestClient(unittest.TestCase):
             groups={},
             person_properties=None,
             group_properties=None,
-            disable_geoip=False,
+            geoip_disable=False,
         )
 
-    @mock.patch("posthog.client.decide")
-    def test_basic_capture_with_feature_flags_switched_off_doesnt_send_them(self, patch_decide):
-        patch_decide.return_value = {"featureFlags": {"beta-feature": "random-variant"}}
+    @mock.patch("posthog.client.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)
@@ -557,13 +713,15 @@ class TestClient(unittest.TestCase):
         self.assertTrue("$feature/beta-feature" not in msg["properties"])
         self.assertTrue("$active_feature_flags" not in msg["properties"])
 
-        self.assertEqual(patch_decide.call_count, 0)
+        self.assertEqual(patch_flags.call_count, 0)
 
     def test_stringifies_distinct_id(self):
         # 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)
@@ -599,7 +757,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
@@ -616,7 +777,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)
@@ -643,7 +807,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)
 
@@ -670,7 +839,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)
@@ -704,7 +876,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")
@@ -724,7 +898,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)
@@ -786,7 +964,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
@@ -796,7 +976,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
@@ -811,7 +993,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)
@@ -882,14 +1066,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)
@@ -914,29 +1102,29 @@ class TestClient(unittest.TestCase):
 
         self.assertEqual(msg, "disabled")
 
-    @mock.patch("posthog.client.decide")
-    def test_disabled_with_feature_flags(self, patch_decide):
+    @mock.patch("posthog.client.flags")
+    def test_disabled_with_feature_flags(self, patch_flags):
         client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, disabled=True)
 
         response = client.get_feature_flag("beta-feature", "12345")
         self.assertIsNone(response)
-        patch_decide.assert_not_called()
+        patch_flags.assert_not_called()
 
         response = client.feature_enabled("beta-feature", "12345")
         self.assertIsNone(response)
-        patch_decide.assert_not_called()
+        patch_flags.assert_not_called()
 
         response = client.get_all_flags("12345")
         self.assertIsNone(response)
-        patch_decide.assert_not_called()
+        patch_flags.assert_not_called()
 
         response = client.get_feature_flag_payload("key", "12345")
         self.assertIsNone(response)
-        patch_decide.assert_not_called()
+        patch_flags.assert_not_called()
 
         response = client.get_all_flags_and_payloads("12345")
         self.assertEqual(response, {"featureFlags": None, "featureFlagPayloads": None})
-        patch_decide.assert_not_called()
+        patch_flags.assert_not_called()
 
         # no capture calls
         self.assertTrue(client.queue.empty())
@@ -970,11 +1158,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)
 
@@ -984,14 +1176,18 @@ class TestClient(unittest.TestCase):
         client.flush()
         self.assertTrue("$geoip_disable" not in msg["properties"])
 
-    @mock.patch("posthog.client.decide")
-    def test_disable_geoip_default_on_decide(self, patch_decide):
-        patch_decide.return_value = {
-            "featureFlags": {"beta-feature": "random-variant", "alpha-feature": True, "off-feature": False}
+    @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,
+            }
         }
         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)
-        patch_decide.assert_called_with(
+        patch_flags.assert_called_with(
             "random_key",
             "https://us.i.posthog.com",
             timeout=3,
@@ -999,11 +1195,13 @@ class TestClient(unittest.TestCase):
             groups={},
             person_properties={"distinct_id": "some_id"},
             group_properties={},
-            disable_geoip=True,
+            geoip_disable=True,
         )
-        patch_decide.reset_mock()
-        client.feature_enabled("random_key", "feature_enabled_distinct_id", disable_geoip=True)
-        patch_decide.assert_called_with(
+        patch_flags.reset_mock()
+        client.feature_enabled(
+            "random_key", "feature_enabled_distinct_id", disable_geoip=True
+        )
+        patch_flags.assert_called_with(
             "random_key",
             "https://us.i.posthog.com",
             timeout=3,
@@ -1011,11 +1209,11 @@ class TestClient(unittest.TestCase):
             groups={},
             person_properties={"distinct_id": "feature_enabled_distinct_id"},
             group_properties={},
-            disable_geoip=True,
+            geoip_disable=True,
         )
-        patch_decide.reset_mock()
+        patch_flags.reset_mock()
         client.get_all_flags_and_payloads("all_flags_payloads_id")
-        patch_decide.assert_called_with(
+        patch_flags.assert_called_with(
             "random_key",
             "https://us.i.posthog.com",
             timeout=3,
@@ -1023,7 +1221,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")
@@ -1038,12 +1236,21 @@ class TestClient(unittest.TestCase):
 
         self.assertFalse(client.feature_enabled("example", "distinct_id"))
 
-    @mock.patch("posthog.client.decide")
-    def test_default_properties_get_added_properly(self, patch_decide):
-        patch_decide.return_value = {
-            "featureFlags": {"beta-feature": "random-variant", "alpha-feature": True, "off-feature": False}
+    @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,
+            }
         }
-        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",
@@ -1051,7 +1258,7 @@ class TestClient(unittest.TestCase):
             person_properties={"x1": "y1"},
             group_properties={"company": {"x": "y"}},
         )
-        patch_decide.assert_called_with(
+        patch_flags.assert_called_with(
             "random_key",
             "http://app2.posthog.com",
             timeout=3,
@@ -1062,10 +1269,10 @@ class TestClient(unittest.TestCase):
                 "company": {"$group_key": "id:5", "x": "y"},
                 "instance": {"$group_key": "app.posthog.com"},
             },
-            disable_geoip=False,
+            geoip_disable=False,
         )
 
-        patch_decide.reset_mock()
+        patch_flags.reset_mock()
         client.get_feature_flag(
             "random_key",
             "some_id",
@@ -1077,7 +1284,7 @@ class TestClient(unittest.TestCase):
                 }
             },
         )
-        patch_decide.assert_called_with(
+        patch_flags.assert_called_with(
             "random_key",
             "http://app2.posthog.com",
             timeout=3,
@@ -1088,13 +1295,15 @@ class TestClient(unittest.TestCase):
                 "company": {"$group_key": "group_override"},
                 "instance": {"$group_key": "app.posthog.com"},
             },
-            disable_geoip=False,
+            geoip_disable=False,
         )
 
-        patch_decide.reset_mock()
+        patch_flags.reset_mock()
         # test nones
-        client.get_all_flags_and_payloads("some_id", groups={}, person_properties=None, group_properties=None)
-        patch_decide.assert_called_with(
+        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",
             timeout=3,
@@ -1102,5 +1311,153 @@ class TestClient(unittest.TestCase):
             groups={},
             person_properties={"distinct_id": "some_id"},
             group_properties={},
-            disable_geoip=False,
+            geoip_disable=False,
         )
+
+    @parameterized.expand(
+        [
+            # name, sys_platform, version_info, expected_runtime, expected_version, expected_os, expected_os_version, platform_method, platform_return, distro_info
+            (
+                "macOS",
+                "darwin",
+                (3, 8, 10),
+                "MockPython",
+                "3.8.10",
+                "Mac OS X",
+                "10.15.7",
+                "mac_ver",
+                ("10.15.7", "", ""),
+                None,
+            ),
+            (
+                "Windows",
+                "win32",
+                (3, 8, 10),
+                "MockPython",
+                "3.8.10",
+                "Windows",
+                "10",
+                "win32_ver",
+                ("10", "", "", ""),
+                None,
+            ),
+            (
+                "Linux",
+                "linux",
+                (3, 8, 10),
+                "MockPython",
+                "3.8.10",
+                "Linux",
+                "20.04",
+                None,
+                None,
+                {"version": "20.04"},
+            ),
+        ]
+    )
+    def test_mock_system_context(
+        self,
+        _name,
+        sys_platform,
+        version_info,
+        expected_runtime,
+        expected_version,
+        expected_os,
+        expected_os_version,
+        platform_method,
+        platform_return,
+        distro_info,
+    ):
+        """Test that we can mock platform and sys for testing system_context"""
+        with mock.patch("posthog.client.platform") as mock_platform:
+            with mock.patch("posthog.client.sys") as mock_sys:
+                # Set up common mocks
+                mock_platform.python_implementation.return_value = expected_runtime
+                mock_sys.version_info = version_info
+                mock_sys.platform = sys_platform
+
+                # Set up platform-specific mocks
+                if platform_method:
+                    getattr(
+                        mock_platform, platform_method
+                    ).return_value = platform_return
+
+                # Special handling for Linux which uses distro module
+                if sys_platform == "linux":
+                    # Directly patch the get_os_info function to return our expected values
+                    with mock.patch(
+                        "posthog.client.get_os_info",
+                        return_value=(expected_os, expected_os_version),
+                    ):
+                        from posthog.client import system_context
+
+                        context = system_context()
+                else:
+                    # Get system context for non-Linux platforms
+                    from posthog.client import system_context
+
+                    context = system_context()
+
+                # Verify results
+                expected_context = {
+                    "$python_runtime": expected_runtime,
+                    "$python_version": expected_version,
+                    "$os": expected_os,
+                    "$os_version": expected_os_version,
+                }
+
+                assert context == expected_context
+
+    @mock.patch("posthog.client.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,
+            },
+            "featureFlagPayloads": {"beta-feature": '{"some": "data"}'},
+            "errorsWhileComputingFlags": False,
+            "requestId": "test-id",
+        }
+
+        client = Client(FAKE_TEST_API_KEY)
+        distinct_id = "test_distinct_id"
+        groups = {"test_group_type": "test_group_id"}
+        person_properties = {"test_property": "test_value"}
+
+        response = client.get_flags_decision(distinct_id, groups, person_properties)
+
+        assert response == {
+            "flags": {
+                "beta-feature": FeatureFlag(
+                    key="beta-feature",
+                    enabled=True,
+                    variant="random-variant",
+                    reason=None,
+                    metadata=LegacyFlagMetadata(
+                        payload='{"some": "data"}',
+                    ),
+                ),
+                "alpha-feature": FeatureFlag(
+                    key="alpha-feature",
+                    enabled=True,
+                    variant=None,
+                    reason=None,
+                    metadata=LegacyFlagMetadata(
+                        payload=None,
+                    ),
+                ),
+                "off-feature": FeatureFlag(
+                    key="off-feature",
+                    enabled=False,
+                    variant=None,
+                    reason=None,
+                    metadata=LegacyFlagMetadata(
+                        payload=None,
+                    ),
+                ),
+            },
+            "errorsWhileComputingFlags": False,
+            "requestId": "test-id",
+        }

posthog/test/test_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

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

posthog/test/test_feature_flag_result.py

@@ -0,0 +1,444 @@
+import unittest
+
+import mock
+
+from posthog.client import Client
+from posthog.test.test_utils import FAKE_TEST_API_KEY
+from posthog.types import FeatureFlag, FeatureFlagResult, FlagMetadata, FlagReason
+
+
+class TestFeatureFlagResult(unittest.TestCase):
+    def test_from_bool_value_and_payload(self):
+        result = FeatureFlagResult.from_value_and_payload(
+            "test-flag", True, "[1, 2, 3]"
+        )
+
+        self.assertEqual(result.key, "test-flag")
+        self.assertEqual(result.enabled, True)
+        self.assertEqual(result.variant, None)
+        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"}'
+        )
+
+        self.assertEqual(result.key, "test-flag")
+        self.assertEqual(result.enabled, False)
+        self.assertEqual(result.variant, None)
+        self.assertEqual(result.payload, {"some": "value"})
+
+    def test_from_variant_value_and_payload(self):
+        result = FeatureFlagResult.from_value_and_payload(
+            "test-flag", "control", "true"
+        )
+
+        self.assertEqual(result.key, "test-flag")
+        self.assertEqual(result.enabled, True)
+        self.assertEqual(result.variant, "control")
+        self.assertEqual(result.payload, True)
+
+    def test_from_none_value_and_payload(self):
+        result = FeatureFlagResult.from_value_and_payload(
+            "test-flag", None, '{"some": "value"}'
+        )
+        self.assertIsNone(result)
+
+    def test_from_boolean_flag_details(self):
+        flag_details = FeatureFlag(
+            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
+            ),
+        )
+
+        result = FeatureFlagResult.from_flag_details(flag_details)
+
+        self.assertEqual(result.key, "test-flag")
+        self.assertEqual(result.enabled, True)
+        self.assertEqual(result.variant, None)
+        self.assertEqual(result.payload, "Some string")
+
+    def test_from_boolean_flag_details_with_override_variant_match_value(self):
+        flag_details = FeatureFlag(
+            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
+            ),
+        )
+
+        result = FeatureFlagResult.from_flag_details(
+            flag_details, override_match_value="control"
+        )
+
+        self.assertEqual(result.key, "test-flag")
+        self.assertEqual(result.enabled, True)
+        self.assertEqual(result.variant, "control")
+        self.assertEqual(result.payload, "Some string")
+
+    def test_from_boolean_flag_details_with_override_boolean_match_value(self):
+        flag_details = FeatureFlag(
+            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
+            ),
+        )
+
+        result = FeatureFlagResult.from_flag_details(
+            flag_details, override_match_value=True
+        )
+
+        self.assertEqual(result.key, "test-flag")
+        self.assertEqual(result.enabled, True)
+        self.assertEqual(result.variant, None)
+        self.assertEqual(result.payload, {"some": "value"})
+
+    def test_from_boolean_flag_details_with_override_false_match_value(self):
+        flag_details = FeatureFlag(
+            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
+            ),
+        )
+
+        result = FeatureFlagResult.from_flag_details(
+            flag_details, override_match_value=False
+        )
+
+        self.assertEqual(result.key, "test-flag")
+        self.assertEqual(result.enabled, False)
+        self.assertEqual(result.variant, None)
+        self.assertEqual(result.payload, {"some": "value"})
+
+    def test_from_variant_flag_details(self):
+        flag_details = FeatureFlag(
+            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
+            ),
+        )
+
+        result = FeatureFlagResult.from_flag_details(flag_details)
+
+        self.assertEqual(result.key, "test-flag")
+        self.assertEqual(result.enabled, True)
+        self.assertEqual(result.variant, "control")
+        self.assertEqual(result.payload, {"some": "value"})
+
+    def test_from_none_flag_details(self):
+        result = FeatureFlagResult.from_flag_details(None)
+
+        self.assertIsNone(result)
+
+    def test_from_flag_details_with_none_payload(self):
+        flag_details = FeatureFlag(
+            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
+            ),
+        )
+
+        result = FeatureFlagResult.from_flag_details(flag_details)
+
+        self.assertEqual(result.key, "test-flag")
+        self.assertEqual(result.enabled, True)
+        self.assertEqual(result.variant, None)
+        self.assertIsNone(result.payload)
+
+
+class TestGetFeatureFlagResult(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        # This ensures no real HTTP POST requests are made
+        cls.capture_patch = mock.patch.object(Client, "capture")
+        cls.capture_patch.start()
+
+    @classmethod
+    def tearDownClass(cls):
+        cls.capture_patch.stop()
+
+    def set_fail(self, e, batch):
+        """Mark the failure handler"""
+        print("FAIL", e, batch)  # noqa: T201
+        self.failed = True
+
+    def setUp(self):
+        self.failed = False
+        self.client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail)
+
+    @mock.patch.object(Client, "capture")
+    def test_get_feature_flag_result_boolean_local_evaluation(self, patch_capture):
+        basic_flag = {
+            "id": 1,
+            "name": "Beta Feature",
+            "key": "person-flag",
+            "active": True,
+            "filters": {
+                "groups": [
+                    {
+                        "properties": [
+                            {
+                                "key": "region",
+                                "operator": "exact",
+                                "value": ["USA"],
+                                "type": "person",
+                            }
+                        ],
+                        "rollout_percentage": 100,
+                    }
+                ],
+                "payloads": {"true": "300"},
+            },
+        }
+        self.client.feature_flags = [basic_flag]
+
+        flag_result = self.client.get_feature_flag_result(
+            "person-flag", "some-distinct-id", person_properties={"region": "USA"}
+        )
+        self.assertEqual(flag_result.enabled, True)
+        self.assertEqual(flag_result.variant, None)
+        self.assertEqual(flag_result.payload, 300)
+        patch_capture.assert_called_with(
+            "some-distinct-id",
+            "$feature_flag_called",
+            {
+                "$feature_flag": "person-flag",
+                "$feature_flag_response": True,
+                "locally_evaluated": True,
+                "$feature/person-flag": True,
+                "$feature_flag_payload": 300,
+            },
+            groups={},
+            disable_geoip=None,
+        )
+
+    @mock.patch.object(Client, "capture")
+    def test_get_feature_flag_result_variant_local_evaluation(self, patch_capture):
+        basic_flag = {
+            "id": 1,
+            "name": "Beta Feature",
+            "key": "person-flag",
+            "active": True,
+            "filters": {
+                "groups": [
+                    {
+                        "properties": [
+                            {
+                                "key": "region",
+                                "operator": "exact",
+                                "value": ["USA"],
+                                "type": "person",
+                            }
+                        ],
+                        "rollout_percentage": 100,
+                    }
+                ],
+                "multivariate": {
+                    "variants": [
+                        {"key": "variant-1", "rollout_percentage": 50},
+                        {"key": "variant-2", "rollout_percentage": 50},
+                    ]
+                },
+                "payloads": {"variant-1": '{"some": "value"}'},
+            },
+        }
+        self.client.feature_flags = [basic_flag]
+
+        flag_result = self.client.get_feature_flag_result(
+            "person-flag", "distinct_id", person_properties={"region": "USA"}
+        )
+        self.assertEqual(flag_result.enabled, True)
+        self.assertEqual(flag_result.variant, "variant-1")
+        self.assertEqual(flag_result.get_value(), "variant-1")
+        self.assertEqual(flag_result.payload, {"some": "value"})
+
+        patch_capture.assert_called_with(
+            "distinct_id",
+            "$feature_flag_called",
+            {
+                "$feature_flag": "person-flag",
+                "$feature_flag_response": "variant-1",
+                "locally_evaluated": True,
+                "$feature/person-flag": "variant-1",
+                "$feature_flag_payload": {"some": "value"},
+            },
+            groups={},
+            disable_geoip=None,
+        )
+
+        another_flag_result = self.client.get_feature_flag_result(
+            "person-flag", "another-distinct-id", person_properties={"region": "USA"}
+        )
+        self.assertEqual(another_flag_result.enabled, True)
+        self.assertEqual(another_flag_result.variant, "variant-2")
+        self.assertEqual(another_flag_result.get_value(), "variant-2")
+        self.assertIsNone(another_flag_result.payload)
+
+        patch_capture.assert_called_with(
+            "another-distinct-id",
+            "$feature_flag_called",
+            {
+                "$feature_flag": "person-flag",
+                "$feature_flag_response": "variant-2",
+                "locally_evaluated": True,
+                "$feature/person-flag": "variant-2",
+            },
+            groups={},
+            disable_geoip=None,
+        )
+
+    @mock.patch("posthog.client.flags")
+    @mock.patch.object(Client, "capture")
+    def test_get_feature_flag_result_boolean_decide(self, patch_capture, patch_flags):
+        patch_flags.return_value = {
+            "flags": {
+                "person-flag": {
+                    "key": "person-flag",
+                    "enabled": True,
+                    "variant": None,
+                    "reason": {
+                        "description": "Matched condition set 1",
+                    },
+                    "metadata": {
+                        "id": 23,
+                        "version": 42,
+                        "payload": "300",
+                    },
+                },
+            },
+        }
+
+        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)
+        patch_capture.assert_called_with(
+            "some-distinct-id",
+            "$feature_flag_called",
+            {
+                "$feature_flag": "person-flag",
+                "$feature_flag_response": True,
+                "locally_evaluated": False,
+                "$feature/person-flag": True,
+                "$feature_flag_reason": "Matched condition set 1",
+                "$feature_flag_id": 23,
+                "$feature_flag_version": 42,
+                "$feature_flag_payload": 300,
+            },
+            groups={},
+            disable_geoip=None,
+        )
+
+    @mock.patch("posthog.client.flags")
+    @mock.patch.object(Client, "capture")
+    def test_get_feature_flag_result_variant_decide(self, patch_capture, patch_flags):
+        patch_flags.return_value = {
+            "flags": {
+                "person-flag": {
+                    "key": "person-flag",
+                    "enabled": True,
+                    "variant": "variant-1",
+                    "reason": {
+                        "description": "Matched condition set 1",
+                    },
+                    "metadata": {
+                        "id": 1,
+                        "version": 2,
+                        "payload": "[1, 2, 3]",
+                    },
+                },
+            },
+        }
+
+        flag_result = self.client.get_feature_flag_result("person-flag", "distinct_id")
+        self.assertEqual(flag_result.enabled, True)
+        self.assertEqual(flag_result.variant, "variant-1")
+        self.assertEqual(flag_result.get_value(), "variant-1")
+        self.assertEqual(flag_result.payload, [1, 2, 3])
+        patch_capture.assert_called_with(
+            "distinct_id",
+            "$feature_flag_called",
+            {
+                "$feature_flag": "person-flag",
+                "$feature_flag_response": "variant-1",
+                "locally_evaluated": False,
+                "$feature/person-flag": "variant-1",
+                "$feature_flag_reason": "Matched condition set 1",
+                "$feature_flag_id": 1,
+                "$feature_flag_version": 2,
+                "$feature_flag_payload": [1, 2, 3],
+            },
+            groups={},
+            disable_geoip=None,
+        )
+
+    @mock.patch("posthog.client.flags")
+    @mock.patch.object(Client, "capture")
+    def test_get_feature_flag_result_unknown_flag(self, patch_capture, patch_flags):
+        patch_flags.return_value = {
+            "flags": {
+                "person-flag": {
+                    "key": "person-flag",
+                    "enabled": True,
+                    "variant": None,
+                    "reason": {
+                        "description": "Matched condition set 1",
+                    },
+                    "metadata": {
+                        "id": 23,
+                        "version": 42,
+                        "payload": "300",
+                    },
+                },
+            },
+        }
+
+        flag_result = self.client.get_feature_flag_result(
+            "no-person-flag", "some-distinct-id"
+        )
+
+        self.assertIsNone(flag_result)
+        patch_capture.assert_called_with(
+            "some-distinct-id",
+            "$feature_flag_called",
+            {
+                "$feature_flag": "no-person-flag",
+                "$feature_flag_response": None,
+                "locally_evaluated": False,
+                "$feature/no-person-flag": None,
+            },
+            groups={},
+            disable_geoip=None,
+        )

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

@@ -2,23 +2,39 @@ import json
 import unittest
 from datetime import date, datetime
 
+import mock
 import pytest
 import requests
 
-from posthog.request import DatetimeSerializer, batch_post, determine_server_host
+from posthog.request import (
+    DatetimeSerializer,
+    QuotaLimitError,
+    batch_post,
+    decide,
+    determine_server_host,
+)
 from posthog.test.test_utils import TEST_API_KEY
 
 
 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)}
@@ -34,16 +50,62 @@ 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):
+        mock_response = requests.Response()
+        mock_response.status_code = 200
+        mock_response._content = json.dumps(
+            {
+                "quotaLimited": ["feature_flags"],
+                "featureFlags": {},
+                "featureFlagPayloads": {},
+                "errorsWhileComputingFlags": False,
+            }
+        ).encode("utf-8")
+
+        with mock.patch("posthog.request._session.post", return_value=mock_response):
+            with self.assertRaises(QuotaLimitError) as cm:
+                decide("fake_key", "fake_host")
+
+            self.assertEqual(cm.exception.status, 200)
+            self.assertEqual(cm.exception.message, "Feature flags quota limited")
+
+    def test_normal_decide_response(self):
+        mock_response = requests.Response()
+        mock_response.status_code = 200
+        mock_response._content = json.dumps(
+            {
+                "featureFlags": {"flag1": True},
+                "featureFlagPayloads": {},
+                "errorsWhileComputingFlags": False,
+            }
+        ).encode("utf-8")
+
+        with mock.patch("posthog.request._session.post", return_value=mock_response):
+            response = decide("fake_key", "fake_host")
+            self.assertEqual(response["featureFlags"], {"flag1": True})
+
 
 @pytest.mark.parametrize(
     "host, expected",

posthog/test/test_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

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

posthog/test/test_utils.py

@@ -1,4 +1,5 @@
 import unittest
+from dataclasses import dataclass
 from datetime import date, datetime, timedelta
 from decimal import Decimal
 from typing import Optional
@@ -6,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
 
@@ -16,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 = {
@@ -53,36 +67,40 @@ 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):
         dict_with_dates = {
             "birthdate": date(1980, 1, 1),
-            "registration": datetime.utcnow(),
+            "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):
@@ -97,34 +115,61 @@ 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
+        class InnerDataClass:
+            inner_foo: str
+            inner_bar: int
+            inner_uuid: UUID
+            inner_date: datetime
+            inner_optional: Optional[str] = None
 
-class TestSizeLimitedDict(unittest.TestCase):
-    def test_size_limited_dict(self):
-        size = 10
-        values = utils.SizeLimitedDict(size, lambda _: -1)
+        @dataclass
+        class TestDataClass:
+            foo: str
+            bar: int
+            nested: InnerDataClass
 
-        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))
+        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,
+            },
+        }

posthog/types.py

@@ -0,0 +1,282 @@
+import json
+from dataclasses import dataclass
+from typing import Any, Callable, List, Optional, TypedDict, Union, cast
+
+FlagValue = Union[bool, str]
+
+# Type alias for the before_send callback function
+# Takes an event dictionary and returns the modified event or None to drop it
+BeforeSendCallback = Callable[[dict[str, Any]], Optional[dict[str, Any]]]
+
+
+@dataclass(frozen=True)
+class FlagReason:
+    code: str
+    condition_index: Optional[int]
+    description: str
+
+    @classmethod
+    def from_json(cls, resp: Any) -> Optional["FlagReason"]:
+        if not resp:
+            return None
+        return cls(
+            code=resp.get("code", ""),
+            condition_index=resp.get("condition_index"),
+            description=resp.get("description", ""),
+        )
+
+
+@dataclass(frozen=True)
+class LegacyFlagMetadata:
+    payload: Any
+
+
+@dataclass(frozen=True)
+class FlagMetadata:
+    id: int
+    payload: Optional[str]
+    version: int
+    description: str
+
+    @classmethod
+    def from_json(cls, resp: Any) -> Union["FlagMetadata", LegacyFlagMetadata]:
+        if not resp:
+            return LegacyFlagMetadata(payload=None)
+        return cls(
+            id=resp.get("id", 0),
+            payload=resp.get("payload"),
+            version=resp.get("version", 0),
+            description=resp.get("description", ""),
+        )
+
+
+@dataclass(frozen=True)
+class FeatureFlag:
+    key: str
+    enabled: bool
+    variant: Optional[str]
+    reason: Optional[FlagReason]
+    metadata: Union[FlagMetadata, LegacyFlagMetadata]
+
+    def get_value(self) -> FlagValue:
+        return self.variant or self.enabled
+
+    @classmethod
+    def from_json(cls, resp: Any) -> "FeatureFlag":
+        reason = None
+        if resp.get("reason"):
+            reason = FlagReason.from_json(resp.get("reason"))
+
+        metadata = None
+        if resp.get("metadata"):
+            metadata = FlagMetadata.from_json(resp.get("metadata"))
+        else:
+            metadata = LegacyFlagMetadata(payload=None)
+
+        return cls(
+            key=resp.get("key"),
+            enabled=resp.get("enabled"),
+            variant=resp.get("variant"),
+            reason=reason,
+            metadata=metadata,
+        )
+
+    @classmethod
+    def from_value_and_payload(
+        cls, key: str, value: FlagValue, payload: Any
+    ) -> "FeatureFlag":
+        enabled, variant = (True, value) if isinstance(value, str) else (value, None)
+        return cls(
+            key=key,
+            enabled=enabled,
+            variant=variant,
+            reason=None,
+            metadata=LegacyFlagMetadata(
+                payload=payload if payload else None,
+            ),
+        )
+
+
+class FlagsResponse(TypedDict, total=False):
+    flags: dict[str, FeatureFlag]
+    errorsWhileComputingFlags: bool
+    requestId: str
+    quotaLimit: Optional[List[str]]
+
+
+class FlagsAndPayloads(TypedDict, total=True):
+    featureFlags: Optional[dict[str, FlagValue]]
+    featureFlagPayloads: Optional[dict[str, Any]]
+
+
+@dataclass(frozen=True)
+class FeatureFlagResult:
+    """
+    The result of calling a feature flag which includes the flag result, variant, and payload.
+
+    Attributes:
+        key (str): The unique identifier of the feature flag.
+        enabled (bool): Whether the feature flag is enabled for the current context.
+        variant (Optional[str]): The variant value if the flag is enabled and has variants, None otherwise.
+        payload (Optional[Any]): Additional data associated with the feature flag, if any.
+        reason (Optional[str]): A description of why the flag was enabled or disabled, if available.
+    """
+
+    key: str
+    enabled: bool
+    variant: Optional[str]
+    payload: Optional[Any]
+    reason: Optional[str]
+
+    def get_value(self) -> FlagValue:
+        """
+        Returns the value of the flag. This is the variant if it exists, otherwise the enabled value.
+        This is the value we report as `$feature_flag_response` in the `$feature_flag_called` event.
+
+        Returns:
+            FlagValue: Either a string variant or boolean value representing the flag's state.
+        """
+        return self.variant or self.enabled
+
+    @classmethod
+    def from_value_and_payload(
+        cls, key: str, value: Union[FlagValue, None], payload: Any
+    ) -> Union["FeatureFlagResult", None]:
+        """
+        Creates a FeatureFlagResult from a flag value and payload.
+
+        Args:
+            key (str): The unique identifier of the feature flag.
+            value (Union[FlagValue, None]): The value of the flag (string variant or boolean).
+            payload (Any): Additional data associated with the feature flag.
+
+        Returns:
+            Union[FeatureFlagResult, None]: A new FeatureFlagResult instance, or None if value is None.
+        """
+        if value is None:
+            return None
+        enabled, variant = (True, value) if isinstance(value, str) else (value, None)
+        return cls(
+            key=key,
+            enabled=enabled,
+            variant=variant,
+            payload=json.loads(payload) if isinstance(payload, str) else payload,
+            reason=None,
+        )
+
+    @classmethod
+    def from_flag_details(
+        cls,
+        details: Union[FeatureFlag, None],
+        override_match_value: Optional[FlagValue] = None,
+    ) -> "FeatureFlagResult | None":
+        """
+        Create a FeatureFlagResult from a FeatureFlag object.
+
+        Args:
+            details (Union[FeatureFlag, None]): The FeatureFlag object to convert.
+            override_match_value (Optional[FlagValue]): If provided, this value will be used to populate
+                the enabled and variant fields instead of the values from the FeatureFlag.
+
+        Returns:
+            FeatureFlagResult | None: A new FeatureFlagResult instance, or None if details is None.
+        """
+
+        if details is None:
+            return None
+
+        if override_match_value is not None:
+            enabled, variant = (
+                (True, override_match_value)
+                if isinstance(override_match_value, str)
+                else (override_match_value, None)
+            )
+        else:
+            enabled, variant = (details.enabled, details.variant)
+
+        return cls(
+            key=details.key,
+            enabled=enabled,
+            variant=variant,
+            payload=(
+                json.loads(details.metadata.payload)
+                if isinstance(details.metadata.payload, str)
+                else details.metadata.payload
+            ),
+            reason=details.reason.description if details.reason else None,
+        )
+
+
+def normalize_flags_response(resp: Any) -> FlagsResponse:
+    """
+    Normalize the response from the decide or flags API endpoint into a FlagsResponse.
+
+    Args:
+        resp: A v3 or v4 response from the decide (or a v1 or v2 response from the flags) API endpoint.
+
+    Returns:
+        A FlagsResponse containing feature flags and their details.
+    """
+    if "requestId" not in resp:
+        resp["requestId"] = None
+    if "flags" in resp:
+        flags = resp["flags"]
+        # For each flag, create a FeatureFlag object
+        for key, value in flags.items():
+            if isinstance(value, FeatureFlag):
+                continue
+            value["key"] = key
+            flags[key] = FeatureFlag.from_json(value)
+    else:
+        # Handle legacy format
+        featureFlags = resp.get("featureFlags", {})
+        featureFlagPayloads = resp.get("featureFlagPayloads", {})
+        resp.pop("featureFlags", None)
+        resp.pop("featureFlagPayloads", None)
+        # look at each key in featureFlags and create a FeatureFlag object
+        flags = {}
+        for key, value in featureFlags.items():
+            flags[key] = FeatureFlag.from_value_and_payload(
+                key, value, featureFlagPayloads.get(key, None)
+            )
+        resp["flags"] = flags
+    return cast(FlagsResponse, resp)
+
+
+def to_flags_and_payloads(resp: FlagsResponse) -> FlagsAndPayloads:
+    """
+    Convert a FlagsResponse into a FlagsAndPayloads object which is a
+    dict of feature flags and their payloads. This is needed by certain
+    functions in the client.
+    Args:
+        resp: A FlagsResponse containing feature flags and their payloads.
+
+    Returns:
+        A tuple containing:
+            - A dictionary mapping flag keys to their values (bool or str)
+            - A dictionary mapping flag keys to their payloads
+    """
+    return {"featureFlags": to_values(resp), "featureFlagPayloads": to_payloads(resp)}
+
+
+def to_values(response: FlagsResponse) -> Optional[dict[str, FlagValue]]:
+    if "flags" not in response:
+        return None
+
+    flags = response.get("flags", {})
+    return {
+        key: value.get_value()
+        for key, value in flags.items()
+        if isinstance(value, FeatureFlag)
+    }
+
+
+def to_payloads(response: FlagsResponse) -> Optional[dict[str, str]]:
+    if "flags" not in response:
+        return None
+
+    return {
+        key: value.metadata.payload
+        for key, value in response.get("flags", {}).items()
+        if isinstance(value, FeatureFlag) and value.enabled and value.metadata.payload
+    }

posthog/utils.py

@@ -2,8 +2,10 @@ import logging
 import numbers
 import re
 from collections import defaultdict
+from dataclasses import asdict, is_dataclass
 from datetime import date, datetime, timezone
 from decimal import Decimal
+from typing import Any, Optional
 from uuid import UUID
 
 import six
@@ -51,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)
@@ -68,6 +72,8 @@ def clean(item):
         pass
     if isinstance(item, dict):
         return _clean_dict(item)
+    if is_dataclass(item) and not isinstance(item, type):
+        return _clean_dataclass(item)
     return _coerce_unicode(item)
 
 
@@ -90,14 +96,41 @@ def _clean_dict(dict_):
     return data
 
 
-def _coerce_unicode(cmplx):
+def _clean_dataclass(dataclass_):
+    data = asdict(dataclass_)
+    data = _clean_dict(data)
+    return data
+
+
+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 = "3.14.1"
+VERSION = "4.6.2"
 
 if __name__ == "__main__":
     print(VERSION, end="")  # noqa: T201

pyproject.toml

@@ -1,10 +1,98 @@
-[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"}
+readme = "README.md"
+requires-python = ">=3.9"
+classifiers = [
+    "Development Status :: 5 - Production/Stable",
+    "Intended Audience :: Developers",
+    "Operating System :: OS Independent",
+    "License :: OSI Approved :: MIT License",
+    "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]
+sentry = ["sentry-sdk", "django"]
+langchain = ["langchain>=0.2.0"]
+dev = [
+    "django-stubs",
+    "lxml",
+    "mypy",
+    "mypy-baseline",
+    "types-mock",
+    "types-python-dateutil",
+    "types-requests",
+    "types-setuptools",
+    "types-six",
+    "pre-commit",
+    "pydantic",
+    "ruff",
+    "setuptools",
+    "packaging", 
+    "wheel",
+    "twine",
+    "tomli",
+    "tomli_w",
+]
+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",
+]
+
+[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",
+]
+
+[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

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

setup.py

@@ -8,94 +8,28 @@ 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",
-    "monotonic>=1.5",
-    "backoff>=1.10.0",
-    "python-dateutil>2.1",
-]
-
-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",
-    ],
-    "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

@@ -1,69 +1,68 @@
 import os
 import sys
+import tomli
+import tomli_w
+import shutil
 
 try:
     from setuptools import setup
 except ImportError:
     from distutils.core import setup
 
-# Don't import module here, since deps may not be installed
+# Don't import analytics-python module here, since deps may not be installed
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), "posthoganalytics"))
-from version import VERSION
+from version import VERSION  # noqa: E402
+
+
+# Copy the original pyproject.toml as backup
+shutil.copy("pyproject.toml", "pyproject.toml.backup")
+
+# Read the original pyproject.toml
+with open("pyproject.toml", "rb") as f:
+    config = tomli.load(f)
+
+# Override specific values
+config["project"]["name"] = "posthoganalytics"
+config["tool"]["setuptools"]["dynamic"]["version"] = {
+    "attr": "posthoganalytics.version.VERSION"
+}
+
+# Rename packages from posthog.* to posthoganalytics.*
+if "packages" in config["tool"]["setuptools"]:
+    new_packages = []
+    for package in config["tool"]["setuptools"]["packages"]:
+        if package == "posthog":
+            new_packages.append("posthoganalytics")
+        elif package.startswith("posthog."):
+            new_packages.append(package.replace("posthog.", "posthoganalytics.", 1))
+        else:
+            new_packages.append(package)
+    config["tool"]["setuptools"]["packages"] = new_packages
+
+# Overwrite the original pyproject.toml
+with open("pyproject.toml", "wb") as f:
+    tomli_w.dump(config, f)
 
 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", "monotonic>=1.5", "backoff>=1.10.0", "python-dateutil>2.1"]
-
-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",
-    ],
     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 :: 2",
-        "Programming Language :: Python :: 2.6",
-        "Programming Language :: Python :: 2.7",
-        "Programming Language :: Python :: 3",
-        "Programming Language :: Python :: 3.2",
-        "Programming Language :: Python :: 3.3",
-        "Programming Language :: Python :: 3.4",
-        "Programming Language :: Python :: 3.5",
-        "Programming Language :: Python :: 3.6",
-        "Programming Language :: Python :: 3.7",
-        "Programming Language :: Python :: 3.8",
-        "Programming Language :: Python :: 3.9",
-        "Programming Language :: Python :: 3.10",
-        "Programming Language :: Python :: 3.11",
-    ],
+    # 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: