remove 'import posthog' (#255)

* chore: more fiddling to get release working

* fix

* fix

.github/workflows/ci.yml

@@ -9,56 +9,61 @@ jobs:
         runs-on: ubuntu-latest
 
         steps:
-            - uses: actions/checkout@v2
+            - uses: actions/checkout@85e6279cec87321a52edac9c87bce653a07cf6c2
               with:
                   fetch-depth: 1
 
-            - name: Set up Python 3.8
-              uses: actions/setup-python@v2
+            - name: Set up Python 3.11
+              uses: actions/setup-python@8d9ed9ac5c53483de85588cdf95a591a75ab9f55
               with:
-                  python-version: 3.8
+                  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 .
+                  ruff format --check .
 
-            - name: Lint with flake8
+            - name: Check types with mypy
               run: |
-                  flake8 posthog --ignore E501,W503
-
-            - name: Check import order with isort
-              run: |
-                  isort --check-only .
+                  mypy --no-site-packages --config-file mypy.ini  . | mypy-baseline filter
 
     tests:
-        name: Python tests
+        name: Python ${{ matrix.python-version }} tests
         runs-on: ubuntu-latest
+        strategy:
+            matrix:
+                python-version: ['3.9', '3.10', '3.11', '3.12', '3.13']
 
         steps:
-            - uses: actions/checkout@v2
+            - uses: actions/checkout@85e6279cec87321a52edac9c87bce653a07cf6c2
               with:
                   fetch-depth: 1
 
-            - name: Set up Python 3.9
-              uses: actions/setup-python@v2
+            - name: Set up Python ${{ matrix.python-version }}
+              uses: actions/setup-python@8d9ed9ac5c53483de85588cdf95a591a75ab9f55
               with:
-                  python-version: 3.9
+                  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,6 +1,186 @@
+## 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
+
+1. Add support for Bedrock Anthropic Usage
+
+## 3.13.0 - 2025-02-12
+
+1. Automatically retry connection errors
+
+## 3.12.1 - 2025-02-11
+
+1. Fix mypy support for 3.12.0
+2. Deprecate `is_simple_flag`
+
 ## 3.12.0 - 2025-02-11
 
 1. Add support for OpenAI beta parse API.
+2. Deprecate `context` parameter
 
 ## 3.11.1 - 2025-02-06
 
@@ -9,6 +189,7 @@
 ## 3.11.0 - 2025-01-28
 
 1. Add the `$ai_span` event to the LangChain callback handler to capture the input and output of intermediary chains.
+
    > LLM observability naming change: event property `$ai_trace_name` is now `$ai_span_name`.
 
 2. Fix serialiazation of Pydantic models in methods.

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,9 +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

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

mypy.ini

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

posthog/__init__.py

@@ -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
@@ -44,13 +54,6 @@ def capture(
     send_feature_flags=False,
     disable_geoip=None,  # type: Optional[bool]
 ):
-    if context is not None:
-        warnings.warn(
-            "The 'context' parameter is deprecated and will be removed in a future version.",
-            DeprecationWarning,
-            stacklevel=2,
-        )
-
     # type: (...) -> Tuple[bool, dict]
     """
     Capture allows you to capture anything a user does within your system, which you can later use in PostHog to find patterns in usage, work out which features to improve or where people are giving up.
@@ -72,6 +75,14 @@ def capture(
     posthog.capture('distinct id', 'purchase', groups={'company': 'id:5'})
     ```
     """
+
+    if context is not None:
+        warnings.warn(
+            "The 'context' parameter is deprecated and will be removed in a future version.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
     return _proxy(
         "capture",
         distinct_id=distinct_id,
@@ -94,13 +105,6 @@ def identify(
     uuid=None,  # type: Optional[str]
     disable_geoip=None,  # type: Optional[bool]
 ):
-    if context is not None:
-        warnings.warn(
-            "The 'context' parameter is deprecated and will be removed in a future version.",
-            DeprecationWarning,
-            stacklevel=2,
-        )
-
     # type: (...) -> Tuple[bool, dict]
     """
     Identify lets you add metadata on your users so you can more easily identify who they are in PostHog, and even do things like segment users by these properties.
@@ -117,6 +121,14 @@ def identify(
     })
     ```
     """
+
+    if context is not None:
+        warnings.warn(
+            "The 'context' parameter is deprecated and will be removed in a future version.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
     return _proxy(
         "identify",
         distinct_id=distinct_id,
@@ -136,13 +148,6 @@ def set(
     uuid=None,  # type: Optional[str]
     disable_geoip=None,  # type: Optional[bool]
 ):
-    if context is not None:
-        warnings.warn(
-            "The 'context' parameter is deprecated and will be removed in a future version.",
-            DeprecationWarning,
-            stacklevel=2,
-        )
-
     # type: (...) -> Tuple[bool, dict]
     """
     Set properties on a user record.
@@ -159,6 +164,14 @@ def set(
      })
      ```
     """
+
+    if context is not None:
+        warnings.warn(
+            "The 'context' parameter is deprecated and will be removed in a future version.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
     return _proxy(
         "set",
         distinct_id=distinct_id,
@@ -178,13 +191,6 @@ def set_once(
     uuid=None,  # type: Optional[str]
     disable_geoip=None,  # type: Optional[bool]
 ):
-    if context is not None:
-        warnings.warn(
-            "The 'context' parameter is deprecated and will be removed in a future version.",
-            DeprecationWarning,
-            stacklevel=2,
-        )
-
     # type: (...) -> Tuple[bool, dict]
     """
     Set properties on a user record, only if they do not yet exist.
@@ -201,6 +207,14 @@ def set_once(
      })
      ```
     """
+
+    if context is not None:
+        warnings.warn(
+            "The 'context' parameter is deprecated and will be removed in a future version.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
     return _proxy(
         "set_once",
         distinct_id=distinct_id,
@@ -221,12 +235,6 @@ def group_identify(
     uuid=None,  # type: Optional[str]
     disable_geoip=None,  # type: Optional[bool]
 ):
-    if context is not None:
-        warnings.warn(
-            "The 'context' parameter is deprecated and will be removed in a future version.",
-            DeprecationWarning,
-            stacklevel=2,
-        )
     # type: (...) -> Tuple[bool, dict]
     """
     Set properties on a group
@@ -243,6 +251,14 @@ def group_identify(
      })
      ```
     """
+
+    if context is not None:
+        warnings.warn(
+            "The 'context' parameter is deprecated and will be removed in a future version.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
     return _proxy(
         "group_identify",
         group_type=group_type,
@@ -263,12 +279,6 @@ def alias(
     uuid=None,  # type: Optional[str]
     disable_geoip=None,  # type: Optional[bool]
 ):
-    if context is not None:
-        warnings.warn(
-            "The 'context' parameter is deprecated and will be removed in a future version.",
-            DeprecationWarning,
-            stacklevel=2,
-        )
     # type: (...) -> Tuple[bool, dict]
     """
     To marry up whatever a user does before they sign up or log in with what they do after you need to make an alias call. This will allow you to answer questions like "Which marketing channels leads to users churning after a month?" or "What do users do on our website before signing up?"
@@ -286,6 +296,14 @@ def alias(
     posthog.alias('anonymous session id', 'distinct id')
     ```
     """
+
+    if context is not None:
+        warnings.warn(
+            "The 'context' parameter is deprecated and will be removed in a future version.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
     return _proxy(
         "alias",
         previous_id=previous_id,
@@ -305,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]
     """
@@ -318,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
@@ -329,6 +349,14 @@ def capture_exception(
 
     ```
     """
+
+    if context is not None:
+        warnings.warn(
+            "The 'context' parameter is deprecated and will be removed in a future version.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
     return _proxy(
         "capture_exception",
         exception=exception,
@@ -338,6 +366,7 @@ def capture_exception(
         timestamp=timestamp,
         uuid=uuid,
         groups=groups,
+        **kwargs,
     )
 
 
@@ -387,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:
@@ -430,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:
@@ -461,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,
@@ -476,6 +505,26 @@ def get_feature_flag_payload(
     )
 
 
+def get_remote_config_payload(
+    key,  # type: str
+):
+    """Get the payload for a remote config feature flag.
+
+    Args:
+        key: The key of the feature flag
+
+    Returns:
+        The payload associated with the feature flag. If payload is encrypted, the return value will decrypted
+
+    Note:
+        Requires personal_api_key to be set for authentication
+    """
+    return _proxy(
+        "get_remote_config_payload",
+        key=key,
+    )
+
+
 def get_all_flags_and_payloads(
     distinct_id,
     groups={},
@@ -483,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,
@@ -531,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(
@@ -553,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,
         )
 
@@ -560,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)
@@ -597,6 +663,9 @@ def _parse_usage_model(
         # Bedrock: https://docs.aws.amazon.com/bedrock/latest/userguide/monitoring-cw.html#runtime-cloudwatch-metrics
         ("inputTokenCount", "input"),
         ("outputTokenCount", "output"),
+        # Bedrock Anthropic
+        ("prompt_tokens", "input"),
+        ("completion_tokens", "output"),
         # langchain-ibm https://pypi.org/project/langchain-ibm/
         ("input_token_count", "input"),
         ("generated_token_count", "output"),
@@ -607,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
@@ -627,9 +698,17 @@ def _parse_usage(response: LLMResult):
 
     if hasattr(response, "generations"):
         for generation in response.generations:
+            if "usage" in generation:
+                llm_usage = _parse_usage_model(generation["usage"])
+                break
+
             for generation_chunk in generation:
-                if generation_chunk.generation_info and ("usage_metadata" in generation_chunk.generation_info):
-                    llm_usage = _parse_usage_model(generation_chunk.generation_info["usage_metadata"])
+                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", {})
@@ -641,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
@@ -663,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,14 +833,18 @@ 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,
     )
 
 
 def _module_in_list(name, items):
-    # type: (str, Optional[List[str]]) -> bool
+    # type: (str | None, Optional[List[str]]) -> bool
     if name is None:
         return False
 
@@ -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

@@ -7,6 +7,8 @@ from typing import Optional
 from dateutil import parser
 from dateutil.relativedelta import relativedelta
 
+from posthog import utils
+from posthog.types import FlagValue
 from posthog.utils import convert_to_datetime_aware, is_valid_regex
 
 __LONG_SCALE__ = float(0xFFFFFFFFFFFFFFF)
@@ -24,7 +26,7 @@ class InconclusiveMatchError(Exception):
 # Given the same distinct_id and key, it'll always return the same float. These floats are
 # uniformly distributed between 0 and 1, so if we want to show this feature to 20% of traffic
 # we can do _hash(key, distinct_id) < 0.2
-def _hash(key, distinct_id, salt=""):
+def _hash(key: str, distinct_id: str, salt: str = "") -> float:
     hash_key = f"{key}.{distinct_id}{salt}"
     hash_val = int(hashlib.sha1(hash_key.encode("utf-8")).hexdigest()[:15], 16)
     return hash_val / __LONG_SCALE__
@@ -41,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
@@ -68,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
@@ -79,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"):
@@ -101,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
@@ -115,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")
@@ -129,8 +149,10 @@ def match_property(property, property_values) -> bool:
 
         def compute_exact_match(value, override_value):
             if isinstance(value, list):
-                return str(override_value).lower() in [str(val).lower() for val in value]
-            return str(value).lower() == str(override_value).lower()
+                return str(override_value).casefold() in [
+                    str(val).casefold() for val in value
+                ]
+            return utils.str_iequals(value, override_value)
 
         if operator == "exact":
             return compute_exact_match(value, override_value)
@@ -141,16 +163,22 @@ def match_property(property, property_values) -> bool:
         return key in property_values
 
     if operator == "icontains":
-        return str(value).lower() in str(override_value).lower()
+        return utils.str_icontains(override_value, value)
 
     if operator == "not_icontains":
-        return str(value).lower() not in str(override_value).lower()
+        return not utils.str_icontains(override_value, value)
 
     if operator == "regex":
-        return is_valid_regex(str(value)) and re.compile(str(value)).search(str(override_value)) is not None
+        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,
@@ -189,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)
@@ -216,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}")
@@ -234,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)
@@ -270,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"
 
@@ -301,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,22 @@ from typing import Any, Optional, Union
 
 import requests
 from dateutil.tz import tzutc
+from urllib3.util.retry import Retry
 
 from posthog.utils import remove_trailing_slash
 from posthog.version import VERSION
 
+# Retry on both connect and read errors
+# by default read errors will only retry idempotent HTTP methods (so not POST)
+adapter = requests.adapters.HTTPAdapter(
+    max_retries=Retry(
+        total=2,
+        connect=2,
+        read=2,
+    )
+)
 _session = requests.sessions.Session()
+_session.mount("https://", adapter)
 
 US_INGESTION_ENDPOINT = "https://us.i.posthog.com"
 EU_INGESTION_ENDPOINT = "https://eu.i.posthog.com"
@@ -32,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")
@@ -41,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"
@@ -66,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)
@@ -75,23 +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 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,
+    )
+
+
 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")
 
 
@@ -105,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,25 +325,33 @@ 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,
             "name": "Beta Feature",
             "key": "beta-feature-local",
-            "is_simple_flag": False,
             "active": True,
             "rollout_percentage": 100,
             "filters": {
                 "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,
                     },
@@ -309,19 +361,33 @@ 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 = {
             "id": 1,
             "name": "Beta Feature",
             "key": "person-flag",
-            "is_simple_flag": True,
             "active": True,
             "filters": {
                 "groups": [
@@ -344,7 +410,6 @@ class TestClient(unittest.TestCase):
             "id": 1,
             "name": "Beta Feature",
             "key": "false-flag",
-            "is_simple_flag": True,
             "active": True,
             "filters": {
                 "groups": [
@@ -369,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 = []
@@ -387,23 +456,50 @@ 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,
             "name": "Beta Feature",
             "key": "beta-feature-local",
-            "is_simple_flag": False,
             "active": True,
             "rollout_percentage": 100,
             "filters": {
                 "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,
                     },
@@ -413,19 +509,33 @@ 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 = {
             "id": 1,
             "name": "Beta Feature",
             "key": "person-flag",
-            "is_simple_flag": True,
             "active": True,
             "filters": {
                 "groups": [
@@ -447,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)
@@ -459,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)
@@ -487,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,
@@ -498,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(
@@ -515,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)
@@ -529,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,
@@ -540,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)
@@ -562,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)
@@ -604,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
@@ -621,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)
@@ -648,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)
 
@@ -675,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)
@@ -709,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")
@@ -729,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)
@@ -791,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
@@ -801,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
@@ -816,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)
@@ -887,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)
@@ -919,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())
@@ -975,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)
 
@@ -989,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,
@@ -1004,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,
@@ -1016,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,
@@ -1028,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")
@@ -1039,16 +1232,25 @@ class TestClient(unittest.TestCase):
 
         patch_get.return_value.raiseError.side_effect = raise_effect
         client = Client(FAKE_TEST_API_KEY, personal_api_key="test")
-        client.feature_flags = [{"key": "example", "is_simple_flag": False}]
+        client.feature_flags = [{"key": "example"}]
 
         self.assertFalse(client.feature_enabled("example", "distinct_id"))
 
-    @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",
@@ -1056,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,
@@ -1067,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",
@@ -1082,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,
@@ -1093,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,
@@ -1107,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
 
 
@@ -125,3 +158,43 @@ def convert_to_datetime_aware(date_obj):
     if date_obj.tzinfo is None:
         date_obj = date_obj.replace(tzinfo=timezone.utc)
     return date_obj
+
+
+def str_icontains(source, search):
+    """
+    Check if a string contains another string, ignoring case.
+
+    Args:
+        source: The string to search within
+        search: The substring to search for
+
+    Returns:
+        bool: True if search is a substring of source (case-insensitive), False otherwise
+
+    Examples:
+        >>> str_icontains("Hello World", "WORLD")
+        True
+        >>> str_icontains("Hello World", "python")
+        False
+    """
+    return str(search).casefold() in str(source).casefold()
+
+
+def str_iequals(value, comparand):
+    """
+    Check if a string equals another string, ignoring case.
+
+    Args:
+        value: The string to compare
+        comparand: The string to compare with
+
+    Returns:
+        bool: True if value and comparand are equal (case-insensitive), False otherwise
+
+    Examples:
+        >>> str_iequals("Hello World", "hello world")
+        True
+        >>> str_iequals("Hello World", "hello")
+        False
+    """
+    return str(value).casefold() == str(comparand).casefold()

posthog/version.py

@@ -1,4 +1,4 @@
-VERSION = "3.12.0"
+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,93 +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",
-        "isort",
-        "flake8",
-        "flake8-print",
-        "pre-commit",
-    ],
-    "test": [
-        "mock>=2.0.0",
-        "freezegun==0.3.15",
-        "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 :: 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
 )

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: