fix: change scoped export, add capturing param (#257)

* change export, add capturing param

* capturing -> capture_exceptions

.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@v1
+            - name: Install uv
+              uses: astral-sh/setup-uv@0c5e2b8115b80b4c7c5ddf6ffdd634974642d182 # v5.4.1
               with:
-                  path: ~/.cache/pip
-                  key: ${{ runner.os }}-pip-${{ hashFiles('setup.py') }}
-                  restore-keys: |
-                      ${{ runner.os }}-pip-
+                  enable-cache: true
+                  pyproject-file: 'pyproject.toml'
 
             - name: Install dev dependencies
+              shell: bash
               run: |
-                  python -m pip install -e .[dev]
-              if: steps.cache.outputs.cache-hit != 'true'
+                  UV_PROJECT_ENVIRONMENT=$pythonLocation uv sync --extra dev
 
-            - name: Check formatting with black
+            - name: Check formatting with ruff
               run: |
-                  black --check .
-            
-            - name: Lint with flake8
-              run: |
-                  flake8 posthog --ignore E501
+                  ruff format --check .
 
-            - name: Check import order with isort
+            - name: Check types with mypy
               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@v1
+            - uses: actions/checkout@85e6279cec87321a52edac9c87bce653a07cf6c2
               with:
                   fetch-depth: 1
 
-            - name: Set up Python 3.7
-              uses: actions/setup-python@v1
+            - name: Set up Python ${{ matrix.python-version }}
+              uses: actions/setup-python@8d9ed9ac5c53483de85588cdf95a591a75ab9f55
               with:
-                  python-version: 3.7
+                  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

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

.pre-commit-config.yaml

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

BEFORE_SEND.md

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

CHANGELOG.md

@@ -1,3 +1,341 @@
+## 4.8.0 - 2025-06-10
+
+- fix: export scoped, rather than tracked, decorator
+- feat: allow use of contexts without error tracking
+
+## 4.7.0 - 2025-06-10
+
+- feat: add support for parse endpoint in responses API (no longer beta)
+
+## 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
+
+1. Fix LangChain callback handler to capture parent run ID.
+
+## 3.11.0 - 2025-01-28
+
+1. Add the `$ai_span` event to the LangChain callback handler to capture the input and output of intermediary chains.
+
+   > LLM observability naming change: event property `$ai_trace_name` is now `$ai_span_name`.
+
+2. Fix serialiazation of Pydantic models in methods.
+
+## 3.10.0 - 2025-01-24
+
+1. Add `$ai_error` and `$ai_is_error` properties to LangChain callback handler, OpenAI, and Anthropic.
+
+## 3.9.3 - 2025-01-23
+
+1. Fix capturing of multiple traces in the LangChain callback handler.
+
+## 3.9.2 - 2025-01-22
+
+1. Fix importing of LangChain callback handler under certain circumstances.
+
+## 3.9.0 - 2025-01-22
+
+1. Add `$ai_trace` event emission to LangChain callback handler.
+
+## 3.8.4 - 2025-01-17
+
+1. Add Anthropic support for LLM Observability.
+2. Update LLM Observability to use output_choices.
+
+## 3.8.3 - 2025-01-14
+
+1. Fix setuptools to include the `posthog.ai.openai` and `posthog.ai.langchain` packages for the `posthoganalytics` package.
+
+## 3.8.2 - 2025-01-14
+
+1. Fix setuptools to include the `posthog.ai.openai` and `posthog.ai.langchain` packages.
+
+## 3.8.1 - 2025-01-14
+
+1. Add LLM Observability with support for OpenAI and Langchain callbacks.
+
+## 3.7.5 - 2025-01-03
+
+1. Add `distinct_id` to group_identify
+
+## 3.7.4 - 2024-11-25
+
+1. Fix bug where this SDK incorrectly sent feature flag events with null values when calling `get_feature_flag_payload`.
+
+## 3.7.3 - 2024-11-25
+
+1. Use personless mode when sending an exception without a provided `distinct_id`.
+
+## 3.7.2 - 2024-11-19
+
+1. Add `type` property to exception stacks.
+
+## 3.7.1 - 2024-10-24
+
+1. Add `platform` property to each frame of exception stacks.
+
+## 3.7.0 - 2024-10-03
+
+1. Adds a new `super_properties` parameter on the client that are appended to every /capture call.
+
+## 3.6.7 - 2024-09-24
+
+1. Remove deprecated datetime.utcnow() in favour of datetime.now(tz=tzutc())
+
+## 3.6.6 - 2024-09-16
+
+1. Fix manual capture support for in app frames
+
+## 3.6.5 - 2024-09-10
+
+1. Fix django integration support for manual exception capture.
+
+## 3.6.4 - 2024-09-05
+
+1. Add manual exception capture.
+
+## 3.6.3 - 2024-09-03
+
+1. Make sure setup.py for posthoganalytics package also discovers the new exception integration package.
+
+## 3.6.2 - 2024-09-03
+
+1. Make sure setup.py discovers the new exception integration package.
+
+## 3.6.1 - 2024-09-03
+
+1. Adds django integration to exception autocapture in alpha state. This feature is not yet stable and may change in future versions.
+
+## 3.6.0 - 2024-08-28
+
+1. Adds exception autocapture in alpha state. This feature is not yet stable and may change in future versions.
+
+## 3.5.2 - 2024-08-21
+
+1. Guard for None values in local evaluation
+
+## 3.5.1 - 2024-08-13
+
+1. Remove "-api" suffix from ingestion hostnames
+
+## 3.5.0 - 2024-02-29
+
+1. - Adds a new `feature_flags_request_timeout_seconds` timeout parameter for feature flags which defaults to 3 seconds, updated from the default 10s for all other API calls.
+
+## 3.4.2 - 2024-02-20
+
+1. Add `historical_migration` option for bulk migration to PostHog Cloud.
+
+## 3.4.1 - 2024-02-09
+
+1. Use new hosts for event capture as well
+
+## 3.4.0 - 2024-02-05
+
+1. Point given hosts to new ingestion hosts
+
+## 3.3.4 - 2024-01-30
+
+1. Update type hints for module variables to work with newer versions of mypy
+
+## 3.3.3 - 2024-01-26
+
+1. Remove new relative date operators, combine into regular date operators
+
+## 3.3.2 - 2024-01-19
+
+1. Return success/failure with all capture calls from module functions
+
+## 3.3.1 - 2024-01-10
+
+1. Make sure we don't override any existing feature flag properties when adding locally evaluated feature flag properties.
+
+## 3.3.0 - 2024-01-09
+
+1. When local evaluation is enabled, we automatically add flag information to all events sent to PostHog, whenever possible. This makes it easier to use these events in experiments.
+
 ## 3.2.0 - 2024-01-09
 
 1. Numeric property handling for feature flags now does the expected: When passed in a number, we do a numeric comparison. When passed in a string, we do a string comparison. Previously, we always did a string comparison.
@@ -15,6 +353,7 @@
 
 1. Restore how feature flags work when the client library is disabled: All requests return `None` and no events are sent when the client is disabled.
 2. Add a `feature_flag_definitions()` debug option, which returns currently loaded feature flag definitions. You can use this to more cleverly decide when to request local evaluation of feature flags.
+
 ## 3.0.0 - 2023-04-14
 
 Breaking change:
@@ -55,7 +394,6 @@ posthog = Posthog('api_key', disable_geoip=False)
 1. Log instead of raise error on posthog personal api key errors
 2. Remove upper bound on backoff dependency
 
-
 ## 2.3.0 - 2023-01-31
 
 1. Add support for returning payloads of matched feature flags
@@ -71,6 +409,7 @@ Changes:
 Changes:
 
 1. Fixes issues with date comparison.
+
 ## 2.1.1 - 2022-09-14
 
 Changes:
@@ -83,6 +422,7 @@ Changes:
 
 1. Feature flag defaults have been removed
 2. Setup logging only when debug mode is enabled.
+
 ## 2.0.1 - 2022-08-04
 
 - Make poll_interval configurable
@@ -94,8 +434,8 @@ Changes:
 Breaking changes:
 
 1. The minimum version requirement for PostHog servers is now 1.38. If you're using PostHog Cloud, you satisfy this requirement automatically.
-2. Feature flag defaults apply only when there's an error fetching feature flag results. Earlier, if the default was set to `True`, even if a flag resolved to `False`, the default would override this. 
-**Note: These are removed in 2.0.2**
+2. Feature flag defaults apply only when there's an error fetching feature flag results. Earlier, if the default was set to `True`, even if a flag resolved to `False`, the default would override this.
+   **Note: These are removed in 2.0.2**
 3. Feature flag remote evaluation doesn't require a personal API key.
 
 New Changes:
@@ -103,18 +443,20 @@ New Changes:
 1. You can now evaluate feature flags locally (i.e. without sending a request to your PostHog servers) by setting a personal API key, and passing in groups and person properties to `is_feature_enabled` and `get_feature_flag` calls.
 2. Introduces a `get_all_flags` method that returns all feature flags. This is useful for when you want to seed your frontend with some initial flags, given a user ID.
 
-
-
 ## 1.4.9 - 2022-06-13
+
 - Support for sending feature flags with capture calls
 
 ## 1.4.8 - 2022-05-12
+
 - Support multi variate feature flags
 
 ## 1.4.7 - 2022-04-25
+
 - Allow feature flags usage without project_api_key
 
 ## 1.4.1 - 2021-05-28
+
 - Fix packaging issues with Sentry integrations
 
 ## 1.4.0 - 2021-05-18

CODEOWNERS

@@ -0,0 +1 @@
+@PostHog/team-feature-flags

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

@@ -1,10 +1,17 @@
 # PostHog Python library example
-
-# Import the library
-import time
+import argparse
 
 import posthog
 
+# Add argument parsing
+parser = argparse.ArgumentParser(description="PostHog Python library example")
+parser.add_argument(
+    "--flag",
+    default="person-on-events-enabled",
+    help="Feature flag key to check (default: person-on-events-enabled)",
+)
+args = parser.parse_args()
+
 posthog.debug = True
 
 # You can find this key on the /setup page in PostHog
@@ -18,7 +25,7 @@ posthog.poll_interval = 10
 
 print(
     posthog.feature_enabled(
-        "person-on-events-enabled",
+        args.flag,  # Use the flag from command line arguments
         "12345",
         groups={"organization": str("0182ee91-8ef7-0000-4cb9-fedc5f00926a")},
         group_properties={
@@ -33,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"))
 
@@ -48,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
@@ -77,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(
@@ -93,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()

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

@@ -1,25 +1,46 @@
 import datetime  # noqa: F401
-from typing import Callable, Dict, Optional  # noqa: F401
+import warnings
+from typing import Callable, Dict, List, Optional, Tuple  # noqa: F401
 
 from posthog.client import Client
+from posthog.exception_capture import Integrations  # noqa: F401
+from posthog.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
+scoped = scoped
+
 """Settings."""
-api_key = None  # type: str
-host = None  # type: str
-on_error = None  # type: Callable
+api_key = None  # type: Optional[str]
+host = None  # type: Optional[str]
+on_error = None  # type: Optional[Callable]
 debug = False  # type: bool
 send = True  # type: bool
 sync_mode = False  # type: bool
 disabled = False  # type: bool
-personal_api_key = None  # type: str
-project_api_key = None  # type: str
+personal_api_key = None  # type: Optional[str]
+project_api_key = None  # type: Optional[str]
 poll_interval = 30  # type: int
 disable_geoip = True  # type: bool
+feature_flags_request_timeout_seconds = 3  # type: int
+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
+privacy_mode = False  # type: bool
 
-default_client = None
+default_client = None  # type: Optional[Client]
 
 
 def capture(
@@ -33,7 +54,7 @@ def capture(
     send_feature_flags=False,
     disable_geoip=None,  # type: Optional[bool]
 ):
-    # type: (...) -> None
+    # 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.
 
@@ -54,7 +75,15 @@ def capture(
     posthog.capture('distinct id', 'purchase', groups={'company': 'id:5'})
     ```
     """
-    _proxy(
+
+    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,
         event=event,
@@ -76,7 +105,7 @@ def identify(
     uuid=None,  # type: Optional[str]
     disable_geoip=None,  # type: Optional[bool]
 ):
-    # type: (...) -> None
+    # 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.
 
@@ -92,7 +121,15 @@ def identify(
     })
     ```
     """
-    _proxy(
+
+    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,
         properties=properties,
@@ -111,7 +148,7 @@ def set(
     uuid=None,  # type: Optional[str]
     disable_geoip=None,  # type: Optional[bool]
 ):
-    # type: (...) -> None
+    # type: (...) -> Tuple[bool, dict]
     """
     Set properties on a user record.
     This will overwrite previous people property values, just like `identify`.
@@ -127,7 +164,15 @@ def set(
      })
      ```
     """
-    _proxy(
+
+    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,
         properties=properties,
@@ -146,7 +191,7 @@ def set_once(
     uuid=None,  # type: Optional[str]
     disable_geoip=None,  # type: Optional[bool]
 ):
-    # type: (...) -> None
+    # type: (...) -> Tuple[bool, dict]
     """
     Set properties on a user record, only if they do not yet exist.
     This will not overwrite previous people property values, unlike `identify`.
@@ -162,7 +207,15 @@ def set_once(
      })
      ```
     """
-    _proxy(
+
+    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,
         properties=properties,
@@ -182,7 +235,7 @@ def group_identify(
     uuid=None,  # type: Optional[str]
     disable_geoip=None,  # type: Optional[bool]
 ):
-    # type: (...) -> None
+    # type: (...) -> Tuple[bool, dict]
     """
     Set properties on a group
 
@@ -198,7 +251,15 @@ def group_identify(
      })
      ```
     """
-    _proxy(
+
+    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,
         group_key=group_key,
@@ -218,7 +279,7 @@ def alias(
     uuid=None,  # type: Optional[str]
     disable_geoip=None,  # type: Optional[bool]
 ):
-    # type: (...) -> None
+    # 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?"
 
@@ -235,7 +296,15 @@ def alias(
     posthog.alias('anonymous session id', 'distinct id')
     ```
     """
-    _proxy(
+
+    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,
         distinct_id=distinct_id,
@@ -246,6 +315,61 @@ def alias(
     )
 
 
+def capture_exception(
+    exception=None,  # type: Optional[BaseException]
+    distinct_id=None,  # type: Optional[str]
+    properties=None,  # type: Optional[Dict]
+    context=None,  # type: Optional[Dict]
+    timestamp=None,  # type: Optional[datetime.datetime]
+    uuid=None,  # type: Optional[str]
+    groups=None,  # type: Optional[Dict]
+    **kwargs,
+):
+    # type: (...) -> Tuple[bool, dict]
+    """
+    capture_exception allows you to capture exceptions that happen in your code. This is useful for debugging and understanding what errors your users are encountering.
+    This function never raises an exception, even if it fails to send the event.
+
+    A `capture_exception` call does not require any fields, but we recommend sending:
+    - `distinct id` which uniquely identifies your user for which this exception happens
+    - `exception` to specify the exception to capture. If not provided, the current exception is captured via `sys.exc_info()`
+
+    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
+    try:
+        1 / 0
+    except Exception as e:
+        posthog.capture_exception(e, 'my specific distinct id')
+        posthog.capture_exception(distinct_id='my specific 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(
+        "capture_exception",
+        exception=exception,
+        distinct_id=distinct_id,
+        properties=properties,
+        context=context,
+        timestamp=timestamp,
+        uuid=uuid,
+        groups=groups,
+        **kwargs,
+    )
+
+
 def feature_enabled(
     key,  # type: str
     distinct_id,  # type: str
@@ -292,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:
@@ -335,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:
@@ -366,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,
@@ -381,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={},
@@ -388,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,
@@ -405,6 +549,11 @@ def feature_flag_definitions():
     return _proxy("feature_flag_definitions")
 
 
+def load_feature_flags():
+    """Load feature flag definitions from PostHog."""
+    return _proxy("load_feature_flags")
+
+
 def page(*args, **kwargs):
     """Send a page call."""
     _proxy("page", *args, **kwargs)
@@ -431,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(
@@ -447,12 +595,25 @@ def _proxy(method, *args, **kwargs):
             poll_interval=poll_interval,
             disabled=disabled,
             disable_geoip=disable_geoip,
+            feature_flags_request_timeout_seconds=feature_flags_request_timeout_seconds,
+            super_properties=super_properties,
+            # TODO: Currently this monitoring begins only when the Client is initialised (which happens when you do something with the SDK)
+            # 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,
         )
 
     # always set incase user changes it
     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

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

posthog/ai/anthropic/anthropic.py

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

posthog/ai/anthropic/anthropic_async.py

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

posthog/ai/anthropic/anthropic_providers.py

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

posthog/ai/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/__init__.py

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

posthog/ai/langchain/callbacks.py

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

posthog/ai/openai/__init__.py

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

posthog/ai/openai/openai.py

@@ -0,0 +1,636 @@
+import time
+import uuid
+from typing import Any, Dict, List, Optional
+
+try:
+    import openai
+except ImportError:
+    raise ModuleNotFoundError(
+        "Please install the OpenAI SDK to use this feature: 'pip install openai'"
+    )
+
+from posthog.ai.utils import (
+    call_llm_and_track_usage,
+    get_model_params,
+    with_privacy_mode,
+)
+from posthog.client import Client as PostHogClient
+
+
+class OpenAI(openai.OpenAI):
+    """
+    A wrapper around the OpenAI SDK that automatically sends LLM usage events to PostHog.
+    """
+
+    _ph_client: PostHogClient
+
+    def __init__(self, posthog_client: PostHogClient, **kwargs):
+        """
+        Args:
+            api_key: OpenAI API key.
+            posthog_client: If provided, events will be captured via this client instead
+                            of the global posthog.
+            **openai_config: Any additional keyword args to set on openai (e.g. organization="xxx").
+        """
+        super().__init__(**kwargs)
+        self._ph_client = posthog_client
+
+        # Store original objects after parent initialization (only if they exist)
+        self._original_chat = getattr(self, "chat", None)
+        self._original_embeddings = getattr(self, "embeddings", None)
+        self._original_beta = getattr(self, "beta", None)
+        self._original_responses = getattr(self, "responses", None)
+
+        # Replace with wrapped versions (only if originals exist)
+        if self._original_chat is not None:
+            self.chat = WrappedChat(self, self._original_chat)
+
+        if self._original_embeddings is not None:
+            self.embeddings = WrappedEmbeddings(self, self._original_embeddings)
+
+        if self._original_beta is not None:
+            self.beta = WrappedBeta(self, self._original_beta)
+
+        if self._original_responses is not None:
+            self.responses = WrappedResponses(self, self._original_responses)
+
+
+class WrappedResponses:
+    """Wrapper for OpenAI responses that tracks usage in PostHog."""
+
+    def __init__(self, client: OpenAI, original_responses):
+        self._client = client
+        self._original = original_responses
+
+    def __getattr__(self, name):
+        """Fallback to original responses object for any methods we don't explicitly handle."""
+        return getattr(self._original, name)
+
+    def create(
+        self,
+        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] = {}
+        final_content = []
+        response = self._original.create(**kwargs)
+
+        def generator():
+            nonlocal usage_stats
+            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 [
+                                "input_tokens",
+                                "output_tokens",
+                                "total_tokens",
+                            ]
+                        }
+
+                        # Add support for cached tokens
+                        if hasattr(chunk.usage, "output_tokens_details") and hasattr(
+                            chunk.usage.output_tokens_details, "reasoning_tokens"
+                        ):
+                            usage_stats["reasoning_tokens"] = (
+                                chunk.usage.output_tokens_details.reasoning_tokens
+                            )
+
+                        if hasattr(chunk.usage, "input_tokens_details") and hasattr(
+                            chunk.usage.input_tokens_details, "cached_tokens"
+                        ):
+                            usage_stats["cache_read_input_tokens"] = (
+                                chunk.usage.input_tokens_details.cached_tokens
+                            )
+
+                    yield chunk
+
+            finally:
+                end_time = time.time()
+                latency = end_time - start_time
+                output = final_content
+                self._capture_streaming_event(
+                    posthog_distinct_id,
+                    posthog_trace_id,
+                    posthog_properties,
+                    posthog_privacy_mode,
+                    posthog_groups,
+                    kwargs,
+                    usage_stats,
+                    latency,
+                    output,
+                )
+
+        return generator()
+
+    def _capture_streaming_event(
+        self,
+        posthog_distinct_id: Optional[str],
+        posthog_trace_id: Optional[str],
+        posthog_properties: Optional[Dict[str, Any]],
+        posthog_privacy_mode: bool,
+        posthog_groups: Optional[Dict[str, Any]],
+        kwargs: Dict[str, Any],
+        usage_stats: Dict[str, int],
+        latency: float,
+        output: Any,
+        tool_calls: Optional[List[Dict[str, Any]]] = None,
+    ):
+        if posthog_trace_id is None:
+            posthog_trace_id = str(uuid.uuid4())
+
+        event_properties = {
+            "$ai_provider": "openai",
+            "$ai_model": kwargs.get("model"),
+            "$ai_model_parameters": get_model_params(kwargs),
+            "$ai_input": with_privacy_mode(
+                self._client._ph_client, posthog_privacy_mode, kwargs.get("input")
+            ),
+            "$ai_output_choices": with_privacy_mode(
+                self._client._ph_client,
+                posthog_privacy_mode,
+                output,
+            ),
+            "$ai_http_status": 200,
+            "$ai_input_tokens": usage_stats.get("input_tokens", 0),
+            "$ai_output_tokens": usage_stats.get("output_tokens", 0),
+            "$ai_cache_read_input_tokens": usage_stats.get(
+                "cache_read_input_tokens", 0
+            ),
+            "$ai_reasoning_tokens": usage_stats.get("reasoning_tokens", 0),
+            "$ai_latency": latency,
+            "$ai_trace_id": posthog_trace_id,
+            "$ai_base_url": str(self._client.base_url),
+            **(posthog_properties or {}),
+        }
+
+        if tool_calls:
+            event_properties["$ai_tools"] = with_privacy_mode(
+                self._client._ph_client,
+                posthog_privacy_mode,
+                tool_calls,
+            )
+
+        if posthog_distinct_id is None:
+            event_properties["$process_person_profile"] = False
+
+        if hasattr(self._client._ph_client, "capture"):
+            self._client._ph_client.capture(
+                distinct_id=posthog_distinct_id or posthog_trace_id,
+                event="$ai_generation",
+                properties=event_properties,
+                groups=posthog_groups,
+            )
+
+    def parse(
+        self,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: bool = False,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        """
+        Parse structured output using OpenAI's 'responses.parse' method, but also track usage in PostHog.
+
+        Args:
+            posthog_distinct_id: Optional ID to associate with the usage event.
+            posthog_trace_id: Optional trace UUID for linking events.
+            posthog_properties: Optional dictionary of extra properties to include in the event.
+            posthog_privacy_mode: Whether to anonymize the input and output.
+            posthog_groups: Optional dictionary of groups to associate with the event.
+            **kwargs: Any additional parameters for the OpenAI Responses Parse API.
+
+        Returns:
+            The response from OpenAI's responses.parse call.
+        """
+        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.parse,
+            **kwargs,
+        )
+
+
+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,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: bool = False,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        """
+        Create an embedding using OpenAI's 'embeddings.create' method, but also track usage in PostHog.
+
+        Args:
+            posthog_distinct_id: Optional ID to associate with the usage event.
+            posthog_trace_id: Optional trace UUID for linking events.
+            posthog_properties: Optional dictionary of extra properties to include in the event.
+            posthog_privacy_mode: Whether to 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 = str(uuid.uuid4())
+
+        start_time = time.time()
+        response = self._original.create(**kwargs)
+        end_time = time.time()
+
+        # Extract usage statistics if available
+        usage_stats = {}
+        if hasattr(response, "usage") and response.usage:
+            usage_stats = {
+                "prompt_tokens": getattr(response.usage, "prompt_tokens", 0),
+                "total_tokens": getattr(response.usage, "total_tokens", 0),
+            }
+
+        latency = end_time - start_time
+
+        # Build the event properties
+        event_properties = {
+            "$ai_provider": "openai",
+            "$ai_model": kwargs.get("model"),
+            "$ai_input": with_privacy_mode(
+                self._client._ph_client, posthog_privacy_mode, kwargs.get("input")
+            ),
+            "$ai_http_status": 200,
+            "$ai_input_tokens": usage_stats.get("prompt_tokens", 0),
+            "$ai_latency": latency,
+            "$ai_trace_id": posthog_trace_id,
+            "$ai_base_url": str(self._client.base_url),
+            **(posthog_properties or {}),
+        }
+
+        if posthog_distinct_id is None:
+            event_properties["$process_person_profile"] = False
+
+        # Send capture event for embeddings
+        if hasattr(self._client._ph_client, "capture"):
+            self._client._ph_client.capture(
+                distinct_id=posthog_distinct_id or posthog_trace_id,
+                event="$ai_embedding",
+                properties=event_properties,
+                groups=posthog_groups,
+            )
+
+        return response
+
+
+class WrappedBeta:
+    """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, self._original.chat)
+
+
+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, self._original.completions)
+
+
+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,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: bool = False,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        return call_llm_and_track_usage(
+            posthog_distinct_id,
+            self._client._ph_client,
+            "openai",
+            posthog_trace_id,
+            posthog_properties,
+            posthog_privacy_mode,
+            posthog_groups,
+            self._client.base_url,
+            self._original.parse,
+            **kwargs,
+        )

posthog/ai/openai/openai_async.py

@@ -0,0 +1,639 @@
+import time
+import uuid
+from typing import Any, Dict, List, Optional
+
+try:
+    import openai
+except ImportError:
+    raise ModuleNotFoundError(
+        "Please install the OpenAI SDK to use this feature: 'pip install openai'"
+    )
+
+from posthog.ai.utils import (
+    call_llm_and_track_usage_async,
+    get_model_params,
+    with_privacy_mode,
+)
+from posthog.client import Client as PostHogClient
+
+
+class AsyncOpenAI(openai.AsyncOpenAI):
+    """
+    An async wrapper around the OpenAI SDK that automatically sends LLM usage events to PostHog.
+    """
+
+    _ph_client: PostHogClient
+
+    def __init__(self, posthog_client: PostHogClient, **kwargs):
+        """
+        Args:
+            api_key: OpenAI API key.
+            posthog_client: If provided, events will be captured via this client 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
+
+        # Store original objects after parent initialization (only if they exist)
+        self._original_chat = getattr(self, "chat", None)
+        self._original_embeddings = getattr(self, "embeddings", None)
+        self._original_beta = getattr(self, "beta", None)
+        self._original_responses = getattr(self, "responses", None)
+
+        # Replace with wrapped versions (only if originals exist)
+        if self._original_chat is not None:
+            self.chat = WrappedChat(self, self._original_chat)
+
+        if self._original_embeddings is not None:
+            self.embeddings = WrappedEmbeddings(self, self._original_embeddings)
+
+        if self._original_beta is not None:
+            self.beta = WrappedBeta(self, self._original_beta)
+
+        if self._original_responses is not None:
+            self.responses = WrappedResponses(self, self._original_responses)
+
+
+class WrappedResponses:
+    """Async wrapper for OpenAI responses that tracks usage in PostHog."""
+
+    def __init__(self, client: AsyncOpenAI, original_responses):
+        self._client = client
+        self._original = original_responses
+
+    def __getattr__(self, name):
+        """Fallback to original responses object for any methods we don't explicitly handle."""
+        return getattr(self._original, name)
+
+    async def create(
+        self,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: bool = False,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        if posthog_trace_id is None:
+            posthog_trace_id = str(uuid.uuid4())
+
+        if kwargs.get("stream", False):
+            return await self._create_streaming(
+                posthog_distinct_id,
+                posthog_trace_id,
+                posthog_properties,
+                posthog_privacy_mode,
+                posthog_groups,
+                **kwargs,
+            )
+
+        return await call_llm_and_track_usage_async(
+            posthog_distinct_id,
+            self._client._ph_client,
+            "openai",
+            posthog_trace_id,
+            posthog_properties,
+            posthog_privacy_mode,
+            posthog_groups,
+            self._client.base_url,
+            self._original.create,
+            **kwargs,
+        )
+
+    async def _create_streaming(
+        self,
+        posthog_distinct_id: Optional[str],
+        posthog_trace_id: Optional[str],
+        posthog_properties: Optional[Dict[str, Any]],
+        posthog_privacy_mode: bool,
+        posthog_groups: Optional[Dict[str, Any]],
+        **kwargs: Any,
+    ):
+        start_time = time.time()
+        usage_stats: Dict[str, int] = {}
+        final_content = []
+        response = await self._original.create(**kwargs)
+
+        async def async_generator():
+            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 [
+                                "input_tokens",
+                                "output_tokens",
+                                "total_tokens",
+                            ]
+                        }
+
+                        # Add support for cached tokens
+                        if hasattr(chunk.usage, "output_tokens_details") and hasattr(
+                            chunk.usage.output_tokens_details, "reasoning_tokens"
+                        ):
+                            usage_stats["reasoning_tokens"] = (
+                                chunk.usage.output_tokens_details.reasoning_tokens
+                            )
+
+                        if hasattr(chunk.usage, "input_tokens_details") and hasattr(
+                            chunk.usage.input_tokens_details, "cached_tokens"
+                        ):
+                            usage_stats["cache_read_input_tokens"] = (
+                                chunk.usage.input_tokens_details.cached_tokens
+                            )
+
+                    yield chunk
+
+            finally:
+                end_time = time.time()
+                latency = end_time - start_time
+                output = final_content
+                await self._capture_streaming_event(
+                    posthog_distinct_id,
+                    posthog_trace_id,
+                    posthog_properties,
+                    posthog_privacy_mode,
+                    posthog_groups,
+                    kwargs,
+                    usage_stats,
+                    latency,
+                    output,
+                )
+
+        return async_generator()
+
+    async def _capture_streaming_event(
+        self,
+        posthog_distinct_id: Optional[str],
+        posthog_trace_id: Optional[str],
+        posthog_properties: Optional[Dict[str, Any]],
+        posthog_privacy_mode: bool,
+        posthog_groups: Optional[Dict[str, Any]],
+        kwargs: Dict[str, Any],
+        usage_stats: Dict[str, int],
+        latency: float,
+        output: Any,
+        tool_calls: Optional[List[Dict[str, Any]]] = None,
+    ):
+        if posthog_trace_id is None:
+            posthog_trace_id = str(uuid.uuid4())
+
+        event_properties = {
+            "$ai_provider": "openai",
+            "$ai_model": kwargs.get("model"),
+            "$ai_model_parameters": get_model_params(kwargs),
+            "$ai_input": with_privacy_mode(
+                self._client._ph_client, posthog_privacy_mode, kwargs.get("input")
+            ),
+            "$ai_output_choices": with_privacy_mode(
+                self._client._ph_client,
+                posthog_privacy_mode,
+                output,
+            ),
+            "$ai_http_status": 200,
+            "$ai_input_tokens": usage_stats.get("input_tokens", 0),
+            "$ai_output_tokens": usage_stats.get("output_tokens", 0),
+            "$ai_cache_read_input_tokens": usage_stats.get(
+                "cache_read_input_tokens", 0
+            ),
+            "$ai_reasoning_tokens": usage_stats.get("reasoning_tokens", 0),
+            "$ai_latency": latency,
+            "$ai_trace_id": posthog_trace_id,
+            "$ai_base_url": str(self._client.base_url),
+            **(posthog_properties or {}),
+        }
+
+        if tool_calls:
+            event_properties["$ai_tools"] = with_privacy_mode(
+                self._client._ph_client,
+                posthog_privacy_mode,
+                tool_calls,
+            )
+
+        if posthog_distinct_id is None:
+            event_properties["$process_person_profile"] = False
+
+        if hasattr(self._client._ph_client, "capture"):
+            self._client._ph_client.capture(
+                distinct_id=posthog_distinct_id or posthog_trace_id,
+                event="$ai_generation",
+                properties=event_properties,
+                groups=posthog_groups,
+            )
+
+    async def parse(
+        self,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: bool = False,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        """
+        Parse structured output using OpenAI's 'responses.parse' method, but also track usage in PostHog.
+
+        Args:
+            posthog_distinct_id: Optional ID to associate with the usage event.
+            posthog_trace_id: Optional trace UUID for linking events.
+            posthog_properties: Optional dictionary of extra properties to include in the event.
+            posthog_privacy_mode: Whether to anonymize the input and output.
+            posthog_groups: Optional dictionary of groups to associate with the event.
+            **kwargs: Any additional parameters for the OpenAI Responses Parse API.
+
+        Returns:
+            The response from OpenAI's responses.parse call.
+        """
+        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,
+            self._original.parse,
+            **kwargs,
+        )
+
+
+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,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: bool = False,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        """
+        Create an embedding using OpenAI's 'embeddings.create' method, but also track usage in PostHog.
+
+        Args:
+            posthog_distinct_id: Optional ID to associate with the usage event.
+            posthog_trace_id: Optional trace UUID for linking events.
+            posthog_properties: Optional dictionary of extra properties to include in the event.
+            posthog_privacy_mode: Whether to 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 = str(uuid.uuid4())
+
+        start_time = time.time()
+        response = await self._original.create(**kwargs)
+        end_time = time.time()
+
+        # Extract usage statistics if available
+        usage_stats = {}
+        if hasattr(response, "usage") and response.usage:
+            usage_stats = {
+                "prompt_tokens": getattr(response.usage, "prompt_tokens", 0),
+                "total_tokens": getattr(response.usage, "total_tokens", 0),
+            }
+
+        latency = end_time - start_time
+
+        # Build the event properties
+        event_properties = {
+            "$ai_provider": "openai",
+            "$ai_model": kwargs.get("model"),
+            "$ai_input": with_privacy_mode(
+                self._client._ph_client, posthog_privacy_mode, kwargs.get("input")
+            ),
+            "$ai_http_status": 200,
+            "$ai_input_tokens": usage_stats.get("prompt_tokens", 0),
+            "$ai_latency": latency,
+            "$ai_trace_id": posthog_trace_id,
+            "$ai_base_url": str(self._client.base_url),
+            **(posthog_properties or {}),
+        }
+
+        if posthog_distinct_id is None:
+            event_properties["$process_person_profile"] = False
+
+        # Send capture event for embeddings
+        if hasattr(self._client._ph_client, "capture"):
+            self._client._ph_client.capture(
+                distinct_id=posthog_distinct_id or posthog_trace_id,
+                event="$ai_embedding",
+                properties=event_properties,
+                groups=posthog_groups,
+            )
+
+        return response
+
+
+class WrappedBeta:
+    """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, self._original.chat)
+
+
+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, self._original.completions)
+
+
+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,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: bool = False,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        return await call_llm_and_track_usage_async(
+            posthog_distinct_id,
+            self._client._ph_client,
+            "openai",
+            posthog_trace_id,
+            posthog_properties,
+            posthog_privacy_mode,
+            posthog_groups,
+            self._client.base_url,
+            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

@@ -0,0 +1,544 @@
+import time
+import uuid
+from typing import Any, Callable, Dict, List, Optional
+
+from httpx import URL
+
+from posthog.client import Client as PostHogClient
+
+
+def get_model_params(kwargs: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Extracts model parameters from the kwargs dictionary.
+    """
+    model_params = {}
+    for param in [
+        "temperature",
+        "max_tokens",  # Deprecated field
+        "max_completion_tokens",
+        "top_p",
+        "frequency_penalty",
+        "presence_penalty",
+        "n",
+        "stop",
+        "stream",  # OpenAI-specific field
+        "streaming",  # Anthropic-specific field
+    ]:
+        if param in kwargs and kwargs[param] is not None:
+            model_params[param] = kwargs[param]
+    return model_params
+
+
+def get_usage(response, provider: str) -> Dict[str, Any]:
+    if provider == "anthropic":
+        return {
+            "input_tokens": response.usage.input_tokens,
+            "output_tokens": response.usage.output_tokens,
+            "cache_read_input_tokens": response.usage.cache_read_input_tokens,
+            "cache_creation_input_tokens": response.usage.cache_creation_input_tokens,
+        }
+    elif provider == "openai":
+        cached_tokens = 0
+        input_tokens = 0
+        output_tokens = 0
+        reasoning_tokens = 0
+
+        # responses api
+        if hasattr(response.usage, "input_tokens"):
+            input_tokens = response.usage.input_tokens
+        if hasattr(response.usage, "output_tokens"):
+            output_tokens = response.usage.output_tokens
+        if hasattr(response.usage, "input_tokens_details") and hasattr(
+            response.usage.input_tokens_details, "cached_tokens"
+        ):
+            cached_tokens = response.usage.input_tokens_details.cached_tokens
+        if hasattr(response.usage, "output_tokens_details") and hasattr(
+            response.usage.output_tokens_details, "reasoning_tokens"
+        ):
+            reasoning_tokens = response.usage.output_tokens_details.reasoning_tokens
+
+        # chat completions
+        if hasattr(response.usage, "prompt_tokens"):
+            input_tokens = response.usage.prompt_tokens
+        if hasattr(response.usage, "completion_tokens"):
+            output_tokens = response.usage.completion_tokens
+        if hasattr(response.usage, "prompt_tokens_details") and hasattr(
+            response.usage.prompt_tokens_details, "cached_tokens"
+        ):
+            cached_tokens = response.usage.prompt_tokens_details.cached_tokens
+
+        return {
+            "input_tokens": input_tokens,
+            "output_tokens": output_tokens,
+            "cache_read_input_tokens": cached_tokens,
+            "reasoning_tokens": reasoning_tokens,
+        }
+    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,
+    }
+
+
+def format_response(response, provider: str):
+    """
+    Format a regular (non-streaming) response.
+    """
+    output = []
+    if response is None:
+        return output
+    if provider == "anthropic":
+        return format_response_anthropic(response)
+    elif provider == "openai":
+        return format_response_openai(response)
+    elif provider == "gemini":
+        return format_response_gemini(response)
+    return output
+
+
+def format_response_anthropic(response):
+    output = []
+    for choice in response.content:
+        if choice.text:
+            output.append(
+                {
+                    "role": "assistant",
+                    "content": choice.text,
+                }
+            )
+    return output
+
+
+def format_response_openai(response):
+    output = []
+    if hasattr(response, "choices"):
+        for choice in response.choices:
+            # Handle Chat Completions response format
+            if hasattr(choice, "message") and choice.message and choice.message.content:
+                output.append(
+                    {
+                        "content": choice.message.content,
+                        "role": choice.message.role,
+                    }
+                )
+    # Handle Responses API format
+    if hasattr(response, "output"):
+        for item in response.output:
+            if item.type == "message":
+                # Extract text content from the content list
+                if hasattr(item, "content") and isinstance(item.content, list):
+                    for content_item in item.content:
+                        if (
+                            hasattr(content_item, "type")
+                            and content_item.type == "output_text"
+                            and hasattr(content_item, "text")
+                        ):
+                            output.append(
+                                {
+                                    "content": content_item.text,
+                                    "role": item.role,
+                                }
+                            )
+                        elif hasattr(content_item, "text"):
+                            output.append(
+                                {
+                                    "content": content_item.text,
+                                    "role": item.role,
+                                }
+                            )
+                        elif (
+                            hasattr(content_item, "type")
+                            and content_item.type == "input_image"
+                            and hasattr(content_item, "image_url")
+                        ):
+                            output.append(
+                                {
+                                    "content": {
+                                        "type": "image",
+                                        "image": content_item.image_url,
+                                    },
+                                    "role": item.role,
+                                }
+                            )
+                else:
+                    output.append(
+                        {
+                            "content": item.content,
+                            "role": item.role,
+                        }
+                    )
+    return output
+
+
+def format_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):
+    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(
+    posthog_distinct_id: Optional[str],
+    ph_client: PostHogClient,
+    provider: str,
+    posthog_trace_id: Optional[str],
+    posthog_properties: Optional[Dict[str, Any]],
+    posthog_privacy_mode: bool,
+    posthog_groups: Optional[Dict[str, Any]],
+    base_url: URL,
+    call_method: Callable[..., Any],
+    **kwargs: Any,
+) -> Any:
+    """
+    Common usage-tracking logic for both sync and async calls.
+    call_method: the llm call method (e.g. openai.chat.completions.create)
+    """
+    start_time = time.time()
+    response = None
+    error = None
+    http_status = 200
+    usage: Dict[str, Any] = {}
+    error_params: Dict[str, any] = {}
+
+    try:
+        response = call_method(**kwargs)
+    except Exception as exc:
+        error = exc
+        http_status = getattr(
+            exc, "status_code", 0
+        )  # default to 0 becuase its likely an SDK error
+        error_params = {
+            "$ai_is_error": True,
+            "$ai_error": exc.__str__(),
+        }
+    finally:
+        end_time = time.time()
+        latency = end_time - start_time
+
+        if posthog_trace_id is None:
+            posthog_trace_id = str(uuid.uuid4())
+
+        if response and (
+            hasattr(response, "usage")
+            or (provider == "gemini" and hasattr(response, "usage_metadata"))
+        ):
+            usage = get_usage(response, provider)
+
+        messages = merge_system_prompt(kwargs, provider)
+
+        event_properties = {
+            "$ai_provider": provider,
+            "$ai_model": kwargs.get("model"),
+            "$ai_model_parameters": get_model_params(kwargs),
+            "$ai_input": with_privacy_mode(ph_client, posthog_privacy_mode, messages),
+            "$ai_output_choices": with_privacy_mode(
+                ph_client, posthog_privacy_mode, format_response(response, provider)
+            ),
+            "$ai_http_status": http_status,
+            "$ai_input_tokens": usage.get("input_tokens", 0),
+            "$ai_output_tokens": usage.get("output_tokens", 0),
+            "$ai_latency": latency,
+            "$ai_trace_id": posthog_trace_id,
+            "$ai_base_url": str(base_url),
+            **(posthog_properties or {}),
+            **(error_params or {}),
+        }
+
+        tool_calls = format_tool_calls(response, provider)
+        if tool_calls:
+            event_properties["$ai_tools"] = with_privacy_mode(
+                ph_client, posthog_privacy_mode, tool_calls
+            )
+
+        if (
+            usage.get("cache_read_input_tokens") is not None
+            and usage.get("cache_read_input_tokens", 0) > 0
+        ):
+            event_properties["$ai_cache_read_input_tokens"] = usage.get(
+                "cache_read_input_tokens", 0
+            )
+
+        if (
+            usage.get("cache_creation_input_tokens") is not None
+            and usage.get("cache_creation_input_tokens", 0) > 0
+        ):
+            event_properties["$ai_cache_creation_input_tokens"] = usage.get(
+                "cache_creation_input_tokens", 0
+            )
+
+        if (
+            usage.get("reasoning_tokens") is not None
+            and usage.get("reasoning_tokens", 0) > 0
+        ):
+            event_properties["$ai_reasoning_tokens"] = usage.get("reasoning_tokens", 0)
+
+        if posthog_distinct_id is None:
+            event_properties["$process_person_profile"] = False
+
+        # Process instructions for Responses API
+        if provider == "openai" and kwargs.get("instructions") is not None:
+            event_properties["$ai_instructions"] = with_privacy_mode(
+                ph_client, posthog_privacy_mode, kwargs.get("instructions")
+            )
+
+        # send the event to posthog
+        if hasattr(ph_client, "capture") and callable(ph_client.capture):
+            ph_client.capture(
+                distinct_id=posthog_distinct_id or posthog_trace_id,
+                event="$ai_generation",
+                properties=event_properties,
+                groups=posthog_groups,
+            )
+
+    if error:
+        raise error
+
+    return response
+
+
+async def call_llm_and_track_usage_async(
+    posthog_distinct_id: Optional[str],
+    ph_client: PostHogClient,
+    provider: str,
+    posthog_trace_id: Optional[str],
+    posthog_properties: Optional[Dict[str, Any]],
+    posthog_privacy_mode: bool,
+    posthog_groups: Optional[Dict[str, Any]],
+    base_url: URL,
+    call_async_method: Callable[..., Any],
+    **kwargs: Any,
+) -> Any:
+    start_time = time.time()
+    response = None
+    error = None
+    http_status = 200
+    usage: Dict[str, Any] = {}
+    error_params: Dict[str, any] = {}
+
+    try:
+        response = await call_async_method(**kwargs)
+    except Exception as exc:
+        error = exc
+        http_status = getattr(
+            exc, "status_code", 0
+        )  # default to 0 because its likely an SDK error
+        error_params = {
+            "$ai_is_error": True,
+            "$ai_error": exc.__str__(),
+        }
+    finally:
+        end_time = time.time()
+        latency = end_time - start_time
+
+        if posthog_trace_id is None:
+            posthog_trace_id = str(uuid.uuid4())
+
+        if response and (
+            hasattr(response, "usage")
+            or (provider == "gemini" and hasattr(response, "usage_metadata"))
+        ):
+            usage = get_usage(response, provider)
+
+        messages = merge_system_prompt(kwargs, provider)
+
+        event_properties = {
+            "$ai_provider": provider,
+            "$ai_model": kwargs.get("model"),
+            "$ai_model_parameters": get_model_params(kwargs),
+            "$ai_input": with_privacy_mode(ph_client, posthog_privacy_mode, messages),
+            "$ai_output_choices": with_privacy_mode(
+                ph_client, posthog_privacy_mode, format_response(response, provider)
+            ),
+            "$ai_http_status": http_status,
+            "$ai_input_tokens": usage.get("input_tokens", 0),
+            "$ai_output_tokens": usage.get("output_tokens", 0),
+            "$ai_latency": latency,
+            "$ai_trace_id": posthog_trace_id,
+            "$ai_base_url": str(base_url),
+            **(posthog_properties or {}),
+            **(error_params or {}),
+        }
+
+        tool_calls = format_tool_calls(response, provider)
+        if tool_calls:
+            event_properties["$ai_tools"] = with_privacy_mode(
+                ph_client, posthog_privacy_mode, tool_calls
+            )
+
+        if (
+            usage.get("cache_read_input_tokens") is not None
+            and usage.get("cache_read_input_tokens", 0) > 0
+        ):
+            event_properties["$ai_cache_read_input_tokens"] = usage.get(
+                "cache_read_input_tokens", 0
+            )
+
+        if (
+            usage.get("cache_creation_input_tokens") is not None
+            and usage.get("cache_creation_input_tokens", 0) > 0
+        ):
+            event_properties["$ai_cache_creation_input_tokens"] = usage.get(
+                "cache_creation_input_tokens", 0
+            )
+
+        if posthog_distinct_id is None:
+            event_properties["$process_person_profile"] = False
+
+        # Process instructions for Responses API
+        if provider == "openai" and kwargs.get("instructions") is not None:
+            event_properties["$ai_instructions"] = with_privacy_mode(
+                ph_client, posthog_privacy_mode, kwargs.get("instructions")
+            )
+
+        # send the event to posthog
+        if hasattr(ph_client, "capture") and callable(ph_client.capture):
+            ph_client.capture(
+                distinct_id=posthog_distinct_id or posthog_trace_id,
+                event="$ai_generation",
+                properties=event_properties,
+                groups=posthog_groups,
+            )
+
+    if error:
+        raise error
+
+    return response
+
+
+def with_privacy_mode(ph_client: PostHogClient, privacy_mode: bool, value: Any):
+    if ph_client.privacy_mode or privacy_mode:
+        return None
+    return value

posthog/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
 
@@ -36,6 +36,7 @@ class Consumer(Thread):
         gzip=False,
         retries=10,
         timeout=15,
+        historical_migration=False,
     ):
         """Create a consumer thread."""
         Thread.__init__(self)
@@ -55,6 +56,7 @@ class Consumer(Thread):
         self.running = True
         self.retries = retries
         self.timeout = timeout
+        self.historical_migration = historical_migration
 
     def run(self):
         """Runs the consumer."""
@@ -94,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
@@ -132,8 +136,17 @@ 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, self.host, gzip=self.gzip, timeout=self.timeout, batch=batch)
+            batch_post(
+                self.api_key,
+                self.host,
+                gzip=self.gzip,
+                timeout=self.timeout,
+                batch=batch,
+                historical_migration=self.historical_migration,
+            )
 
         send_request()

posthog/exception_capture.py

@@ -0,0 +1,71 @@
+# 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
+from enum import Enum
+from typing import TYPE_CHECKING, List, Optional
+
+if TYPE_CHECKING:
+    from posthog.client import Client
+
+
+class Integrations(str, Enum):
+    Django = "django"
+
+
+class ExceptionCapture:
+    # TODO: Add client side rate limiting to prevent spamming the server with exceptions
+
+    log = logging.getLogger("posthog")
+
+    def __init__(
+        self, client: "Client", integrations: Optional[List[Integrations]] = None
+    ):
+        self.client = client
+        self.original_excepthook = sys.excepthook
+        sys.excepthook = self.exception_handler
+        threading.excepthook = self.thread_exception_handler
+        self.enabled_integrations = []
+
+        for integration in integrations or []:
+            # TODO: Maybe find a better way of enabling integrations
+            # This is very annoying currently if we had to add any configuration per integration
+            if integration == Integrations.Django:
+                try:
+                    from posthog.exception_integrations.django import DjangoIntegration
+
+                    enabled_integration = DjangoIntegration(self.exception_receiver)
+                    self.enabled_integrations.append(enabled_integration)
+                except Exception as e:
+                    self.log.exception(f"Failed to enable Django integration: {e}")
+
+    def close(self):
+        sys.excepthook = self.original_excepthook
+        for integration in self.enabled_integrations:
+            integration.uninstall()
+
+    def exception_handler(self, exc_type, exc_value, exc_traceback):
+        # don't affect default behaviour.
+        self.capture_exception((exc_type, exc_value, exc_traceback))
+        self.original_excepthook(exc_type, exc_value, exc_traceback)
+
+    def thread_exception_handler(self, args):
+        self.capture_exception((args.exc_type, args.exc_value, args.exc_traceback))
+
+    def exception_receiver(self, exc_info, extra_properties):
+        if "distinct_id" in extra_properties:
+            metadata = {"distinct_id": extra_properties["distinct_id"]}
+        else:
+            metadata = None
+        self.capture_exception((exc_info[0], exc_info[1], exc_info[2]), metadata)
+
+    def capture_exception(self, exception, metadata=None):
+        try:
+            distinct_id = metadata.get("distinct_id") if metadata else None
+            self.client.capture_exception(exception, distinct_id)
+        except Exception as e:
+            self.log.exception(f"Failed to capture exception: {e}")

posthog/exception_integrations/__init__.py

@@ -0,0 +1,5 @@
+class IntegrationEnablingError(Exception):
+    """
+    The integration could not be enabled due to a user error like
+    `django` not being installed for the `DjangoIntegration`.
+    """

posthog/exception_integrations/django.py

@@ -0,0 +1,117 @@
+# 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
+
+from posthog.exception_integrations import IntegrationEnablingError
+
+try:
+    from django import VERSION as DJANGO_VERSION
+    from django.core import signals
+
+except ImportError:
+    raise IntegrationEnablingError("Django not installed")
+
+
+if TYPE_CHECKING:
+    from typing import Any, Dict  # noqa: F401
+
+    from django.core.handlers.wsgi import WSGIRequest  # noqa: F401
+
+
+class DjangoIntegration:
+    # TODO: Abstract integrations one we have more and can see patterns
+    """
+    Autocapture errors from a Django application.
+    """
+
+    identifier = "django"
+
+    def __init__(self, capture_exception_fn=None):
+        if DJANGO_VERSION < (4, 2):
+            raise IntegrationEnablingError("Django 4.2 or newer is required.")
+
+        # TODO: Right now this seems too complicated / overkill for us, but seems like we can automatically plug in middlewares
+        # which is great for users (they don't need to do this) and everything should just work.
+        # We should consider this in the future, but for now we can just use the middleware and signals handlers.
+        # See: https://github.com/getsentry/sentry-python/blob/269d96d6e9821122fbff280e6a26956e5ed03c0b/sentry_sdk/integrations/django/__init__.py
+
+        self.capture_exception_fn = capture_exception_fn
+
+        def _got_request_exception(request=None, **kwargs):
+            # type: (WSGIRequest, **Any) -> None
+
+            extra_props = {}
+            if request is not None:
+                # get headers metadata
+                extra_props = DjangoRequestExtractor(request).extract_person_data()
+
+            self.capture_exception_fn(sys.exc_info(), extra_props)
+
+        signals.got_request_exception.connect(_got_request_exception)
+
+    def uninstall(self):
+        pass
+
+
+class DjangoRequestExtractor:
+    def __init__(self, request):
+        # type: (Any) -> None
+        self.request = request
+
+    def extract_person_data(self):
+        headers = self.headers()
+
+        # Extract traceparent and tracestate headers
+        traceparent = headers.get("Traceparent")
+        tracestate = headers.get("Tracestate")
+
+        # Extract the distinct_id from tracestate
+        distinct_id = None
+        if tracestate:
+            # TODO: Align on the format of the distinct_id in tracestate
+            # We can't have comma or equals in header values here, so maybe we should base64 encode it?
+            match = re.search(r"posthog-distinct-id=([^,]+)", tracestate)
+            if match:
+                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

@@ -0,0 +1,926 @@
+# 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
+
+import linecache
+import os
+import re
+import sys
+from datetime import datetime
+from typing import TYPE_CHECKING
+
+try:
+    # Python 3.11
+    from builtins import BaseExceptionGroup
+except ImportError:
+    # Python 3.10 and below
+    BaseExceptionGroup = None  # type: ignore
+
+
+DEFAULT_MAX_VALUE_LENGTH = 1024
+
+
+if TYPE_CHECKING:
+    from types import FrameType, TracebackType
+    from typing import (  # noqa: F401
+        Any,
+        Callable,
+        Dict,
+        Iterator,
+        List,
+        Literal,
+        Optional,
+        Set,
+        Tuple,
+        Type,
+        TypedDict,
+        TypeVar,
+        Union,
+        cast,
+    )
+
+    ExcInfo = Union[
+        Tuple[Type[BaseException], BaseException, Optional[TracebackType]],
+        Tuple[None, None, None],
+    ]
+    LogLevelStr = Literal["fatal", "critical", "error", "warning", "info", "debug"]
+
+    Event = TypedDict(
+        "Event",
+        {
+            "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,
+            "duration": Optional[float],
+            "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
+            # "extra": MutableMapping[str, object],
+            # "fingerprint": List[str],
+            "level": LogLevelStr,
+            # "logentry": Mapping[str, object],
+            "logger": str,
+            # "measurements": Dict[str, MeasurementValue],
+            "message": str,
+            "modules": Dict[str, str],
+            # "monitor_config": Mapping[str, object],
+            "monitor_slug": Optional[str],
+            "platform": Literal["python"],
+            "profile": object,  # Should be sentry_sdk.profiler.Profile, but we can't import that here due to circular imports
+            "release": str,
+            "request": Dict[str, object],
+            # "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
+            "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
+            "timestamp": Optional[datetime],  # Must be set before sending the event
+            "transaction": str,
+            # "transaction_info": Mapping[str, Any],  # TODO: We can expand on this type
+            "type": Literal["check_in", "transaction"],
+            "user": Dict[str, object],
+            "_metrics_summary": Dict[str, object],
+        },
+        total=False,
+    )
+
+
+epoch = datetime(1970, 1, 1)
+
+
+BASE64_ALPHABET = re.compile(r"^[a-zA-Z0-9/+=]*$")
+
+SENSITIVE_DATA_SUBSTITUTE = "[Filtered]"
+
+
+def to_timestamp(value):
+    # type: (datetime) -> float
+    return (value - epoch).total_seconds()
+
+
+def format_timestamp(value):
+    # type: (datetime) -> str
+    return value.strftime("%Y-%m-%dT%H:%M:%S.%fZ")
+
+
+def event_hint_with_exc_info(exc_info=None):
+    # type: (Optional[ExcInfo]) -> Dict[str, Optional[ExcInfo]]
+    """Creates a hint with the exc info filled in."""
+    if exc_info is None:
+        exc_info = sys.exc_info()
+    else:
+        exc_info = exc_info_from_error(exc_info)
+    if exc_info[0] is None:
+        exc_info = None
+    return {"exc_info": exc_info}
+
+
+class AnnotatedValue:
+    """
+    Meta information for a data field in the event payload.
+    This is to tell Relay that we have tampered with the fields value.
+    See:
+    https://github.com/getsentry/relay/blob/be12cd49a0f06ea932ed9b9f93a655de5d6ad6d1/relay-general/src/types/meta.rs#L407-L423
+    """
+
+    __slots__ = ("value", "metadata")
+
+    def __init__(self, value, metadata):
+        # type: (Optional[Any], Dict[str, Any]) -> None
+        self.value = value
+        self.metadata = metadata
+
+    def __eq__(self, other):
+        # type: (Any) -> bool
+        if not isinstance(other, AnnotatedValue):
+            return False
+
+        return self.value == other.value and self.metadata == other.metadata
+
+    @classmethod
+    def removed_because_raw_data(cls):
+        # type: () -> AnnotatedValue
+        """The value was removed because it could not be parsed. This is done for request body values that are not json nor a form."""
+        return AnnotatedValue(
+            value="",
+            metadata={
+                "rem": [  # Remark
+                    [
+                        "!raw",  # Unparsable raw data
+                        "x",  # The fields original value was removed
+                    ]
+                ]
+            },
+        )
+
+    @classmethod
+    def removed_because_over_size_limit(cls):
+        # type: () -> AnnotatedValue
+        """The actual value was removed because the size of the field exceeded the configured maximum size (specified with the max_request_body_size sdk option)"""
+        return AnnotatedValue(
+            value="",
+            metadata={
+                "rem": [  # Remark
+                    [
+                        "!config",  # Because of configured maximum size
+                        "x",  # The fields original value was removed
+                    ]
+                ]
+            },
+        )
+
+    @classmethod
+    def substituted_because_contains_sensitive_data(cls):
+        # type: () -> AnnotatedValue
+        """The actual value was removed because it contained sensitive information."""
+        return AnnotatedValue(
+            value=SENSITIVE_DATA_SUBSTITUTE,
+            metadata={
+                "rem": [  # Remark
+                    [
+                        "!config",  # Because of SDK configuration (in this case the config is the hard coded removal of certain django cookies)
+                        "s",  # The fields original value was substituted
+                    ]
+                ]
+            },
+        )
+
+
+if TYPE_CHECKING:
+    T = TypeVar("T")
+    Annotated = Union[AnnotatedValue, T]
+
+
+def get_type_name(cls):
+    # type: (Optional[type]) -> Optional[str]
+    return getattr(cls, "__qualname__", None) or getattr(cls, "__name__", None)
+
+
+def get_type_module(cls):
+    # type: (Optional[type]) -> Optional[str]
+    mod = getattr(cls, "__module__", None)
+    if mod not in (None, "builtins", "__builtins__"):
+        return mod
+    return None
+
+
+def should_hide_frame(frame: "FrameType") -> bool:
+    try:
+        mod = frame.f_globals["__name__"]
+        if mod.startswith("sentry_sdk."):
+            return True
+    except (AttributeError, KeyError):
+        pass
+
+    for flag_name in "__traceback_hide__", "__tracebackhide__":
+        try:
+            if frame.f_locals[flag_name]:
+                return True
+        except Exception:
+            pass
+
+    return False
+
+
+def iter_stacks(tb):
+    # type: (Optional[TracebackType]) -> Iterator[TracebackType]
+    tb_ = tb  # type: Optional[TracebackType]
+    while tb_ is not None:
+        if not should_hide_frame(tb_.tb_frame):
+            yield tb_
+        tb_ = tb_.tb_next
+
+
+def get_lines_from_file(
+    filename,  # type: str
+    lineno,  # type: int
+    max_length=None,  # type: Optional[int]
+    loader=None,  # type: Optional[Any]
+    module=None,  # type: Optional[str]
+):
+    # type: (...) -> Tuple[List[Annotated[str]], Optional[Annotated[str]], List[Annotated[str]]]
+    context_lines = 5
+    source = None
+    if loader is not None and hasattr(loader, "get_source"):
+        try:
+            source_str = loader.get_source(module)  # type: Optional[str]
+        except (ImportError, IOError):
+            source_str = None
+        if source_str is not None:
+            source = source_str.splitlines()
+
+    if source is None:
+        try:
+            source = linecache.getlines(filename)
+        except (OSError, IOError):
+            return [], None, []
+
+    if not source:
+        return [], None, []
+
+    lower_bound = max(0, lineno - context_lines)
+    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]
+        ]
+        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)
+            for line in source[(lineno + 1) : upper_bound]  # noqa: E203
+        ]
+        return pre_context, context_line, post_context
+    except IndexError:
+        # the file may have changed since it was loaded into memory
+        return [], None, []
+
+
+def get_source_context(
+    frame,  # type: FrameType
+    tb_lineno,  # type: int
+    max_value_length=None,  # type: Optional[int]
+):
+    # type: (...) -> Tuple[List[Annotated[str]], Optional[Annotated[str]], List[Annotated[str]]]
+    try:
+        abs_path = frame.f_code.co_filename  # type: Optional[str]
+    except Exception:
+        abs_path = None
+    try:
+        module = frame.f_globals["__name__"]
+    except Exception:
+        return [], None, []
+    try:
+        loader = frame.f_globals["__loader__"]
+    except Exception:
+        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 [], None, []
+
+
+def safe_str(value):
+    # type: (Any) -> str
+    try:
+        return str(value)
+    except Exception:
+        return safe_repr(value)
+
+
+def safe_repr(value):
+    # type: (Any) -> str
+    try:
+        return repr(value)
+    except Exception:
+        return "<broken repr>"
+
+
+def filename_for_module(module, abs_path):
+    # type: (Optional[str], Optional[str]) -> Optional[str]
+    if not abs_path or not module:
+        return abs_path
+
+    try:
+        if abs_path.endswith(".pyc"):
+            abs_path = abs_path[:-1]
+
+        base_module = module.split(".", 1)[0]
+        if base_module == module:
+            return os.path.basename(abs_path)
+
+        base_module_path = sys.modules[base_module].__file__
+        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
+        )
+    except Exception:
+        return abs_path
+
+
+def serialize_frame(
+    frame,
+    tb_lineno=None,
+    include_local_variables=True,
+    include_source_context=True,
+    max_value_length=None,
+    custom_repr=None,
+):
+    # type: (FrameType, Optional[int], bool, bool, Optional[int], Optional[Callable[..., Optional[str]]]) -> Dict[str, Any]
+    f_code = getattr(frame, "f_code", None)
+    if not f_code:
+        abs_path = None
+        function = None
+    else:
+        abs_path = frame.f_code.co_filename
+        function = frame.f_code.co_name
+    try:
+        module = frame.f_globals["__name__"]
+    except Exception:
+        module = None
+
+    if tb_lineno is None:
+        tb_lineno = frame.f_lineno
+
+    rv = {
+        "platform": "python",
+        "filename": filename_for_module(module, abs_path) or None,
+        "abs_path": os.path.abspath(abs_path) if abs_path else None,
+        "function": function or "<unknown>",
+        "module": module,
+        "lineno": tb_lineno,
+    }  # type: Dict[str, Any]
+
+    if include_source_context:
+        rv["pre_context"], rv["context_line"], rv["post_context"] = get_source_context(
+            frame, tb_lineno, max_value_length
+        )
+
+    if include_local_variables:
+        # TODO(nk): Sort out this current invalid import
+        # from sentry_sdk.serializer import serialize
+
+        # rv["vars"] = serialize(
+        #     dict(frame.f_locals), is_vars=True, custom_repr=custom_repr
+        # )
+        pass
+
+    return rv
+
+
+def current_stacktrace(
+    include_local_variables=True,  # type: bool
+    include_source_context=True,  # type: bool
+    max_value_length=None,  # type: Optional[int]
+):
+    # type: (...) -> Dict[str, Any]
+    __tracebackhide__ = True
+    frames = []
+
+    f = sys._getframe()  # type: Optional[FrameType]
+    while f is not None:
+        if not should_hide_frame(f):
+            frames.append(
+                serialize_frame(
+                    f,
+                    include_local_variables=include_local_variables,
+                    include_source_context=include_source_context,
+                    max_value_length=max_value_length,
+                )
+            )
+        f = f.f_back
+
+    frames.reverse()
+
+    return {"frames": frames, "type": "raw"}
+
+
+def get_errno(exc_value):
+    # type: (BaseException) -> Optional[Any]
+    return getattr(exc_value, "errno", None)
+
+
+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)
+    )
+
+
+def single_exception_from_error_tuple(
+    exc_type,  # type: Optional[type]
+    exc_value,  # type: Optional[BaseException]
+    tb,  # type: Optional[TracebackType]
+    client_options=None,  # type: Optional[Dict[str, Any]]
+    mechanism=None,  # type: Optional[Dict[str, Any]]
+    exception_id=None,  # type: Optional[int]
+    parent_id=None,  # type: Optional[int]
+    source=None,  # type: Optional[str]
+):
+    # type: (...) -> Dict[str, Any]
+    """
+    Creates a dict that goes into the events `exception.values` list and is ingestible by Sentry.
+
+    See the Exception Interface documentation for more details:
+    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}
+    )
+    if exception_id is not None:
+        exception_value["mechanism"]["exception_id"] = exception_id
+
+    if exc_value is not None:
+        errno = get_errno(exc_value)
+    else:
+        errno = None
+
+    if errno is not None:
+        exception_value["mechanism"].setdefault("meta", {}).setdefault(
+            "errno", {}
+        ).setdefault("number", errno)
+
+    if source is not None:
+        exception_value["mechanism"]["source"] = source
+
+    is_root_exception = exception_id == 0
+    if not is_root_exception and parent_id is not None:
+        exception_value["mechanism"]["parent_id"] = parent_id
+        exception_value["mechanism"]["type"] = "chained"
+
+    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
+    )
+    if is_exception_group:
+        exception_value["mechanism"]["is_exception_group"] = True
+
+    exception_value["module"] = get_type_module(exc_type)
+    exception_value["type"] = get_type_name(exc_type)
+    exception_value["value"] = get_error_message(exc_value)
+
+    if client_options is None:
+        include_local_variables = True
+        include_source_context = True
+        max_value_length = DEFAULT_MAX_VALUE_LENGTH  # fallback
+        custom_repr = None
+    else:
+        include_local_variables = client_options["include_local_variables"]
+        include_source_context = client_options["include_source_context"]
+        max_value_length = client_options["max_value_length"]
+        custom_repr = client_options.get("custom_repr")
+
+    frames = [
+        serialize_frame(
+            tb.tb_frame,
+            tb_lineno=tb.tb_lineno,
+            include_local_variables=include_local_variables,
+            include_source_context=include_source_context,
+            max_value_length=max_value_length,
+            custom_repr=custom_repr,
+        )
+        for tb in iter_stacks(tb)
+    ]
+
+    if frames:
+        exception_value["stacktrace"] = {"frames": frames, "type": "raw"}
+
+    return exception_value
+
+
+HAS_CHAINED_EXCEPTIONS = hasattr(Exception, "__suppress_context__")
+
+if HAS_CHAINED_EXCEPTIONS:
+
+    def walk_exception_chain(exc_info):
+        # type: (ExcInfo) -> Iterator[ExcInfo]
+        exc_type, exc_value, tb = exc_info
+
+        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
+        ):
+            yield exc_type, exc_value, tb
+
+            # Avoid hashing random types we don't know anything
+            # about. Use the list to keep a ref so that the `id` is
+            # not used for another object.
+            seen_exceptions.append(exc_value)
+            seen_exception_ids.add(id(exc_value))
+
+            if exc_value.__suppress_context__:
+                cause = exc_value.__cause__
+            else:
+                cause = exc_value.__context__
+            if cause is None:
+                break
+            exc_type = type(cause)
+            exc_value = cause
+            tb = getattr(cause, "__traceback__", None)
+
+else:
+
+    def walk_exception_chain(exc_info):
+        # type: (ExcInfo) -> Iterator[ExcInfo]
+        yield exc_info
+
+
+def exceptions_from_error(
+    exc_type,  # type: Optional[type]
+    exc_value,  # type: Optional[BaseException]
+    tb,  # type: Optional[TracebackType]
+    client_options=None,  # type: Optional[Dict[str, Any]]
+    mechanism=None,  # type: Optional[Dict[str, Any]]
+    exception_id=0,  # type: int
+    parent_id=0,  # type: int
+    source=None,  # type: Optional[str]
+):
+    # type: (...) -> Tuple[int, List[Dict[str, Any]]]
+    """
+    Creates the list of exceptions.
+    This can include chained exceptions and exceptions from an ExceptionGroup.
+
+    See the Exception Interface documentation for more details:
+    https://develop.sentry.dev/sdk/event-payloads/exception/
+    """
+
+    parent = single_exception_from_error_tuple(
+        exc_type=exc_type,
+        exc_value=exc_value,
+        tb=tb,
+        client_options=client_options,
+        mechanism=mechanism,
+        exception_id=exception_id,
+        parent_id=parent_id,
+        source=source,
+    )
+    exceptions = [parent]
+
+    parent_id = exception_id
+    exception_id += 1
+
+    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
+        )
+        if exception_has_cause:
+            cause = exc_value.__cause__  # type: ignore
+            (exception_id, child_exceptions) = exceptions_from_error(
+                exc_type=type(cause),
+                exc_value=cause,
+                tb=getattr(cause, "__traceback__", None),
+                client_options=client_options,
+                mechanism=mechanism,
+                exception_id=exception_id,
+                source="__cause__",
+            )
+            exceptions.extend(child_exceptions)
+
+    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
+        )
+        if exception_has_content:
+            context = exc_value.__context__  # type: ignore
+            (exception_id, child_exceptions) = exceptions_from_error(
+                exc_type=type(context),
+                exc_value=context,
+                tb=getattr(context, "__traceback__", None),
+                client_options=client_options,
+                mechanism=mechanism,
+                exception_id=exception_id,
+                source="__context__",
+            )
+            exceptions.extend(child_exceptions)
+
+    # Add exceptions from an ExceptionGroup.
+    is_exception_group = exc_value and hasattr(exc_value, "exceptions")
+    if is_exception_group:
+        for idx, e in enumerate(exc_value.exceptions):  # type: ignore
+            (exception_id, child_exceptions) = exceptions_from_error(
+                exc_type=type(e),
+                exc_value=e,
+                tb=getattr(e, "__traceback__", None),
+                client_options=client_options,
+                mechanism=mechanism,
+                exception_id=exception_id,
+                parent_id=parent_id,
+                source="exceptions[%s]" % idx,
+            )
+            exceptions.extend(child_exceptions)
+
+    return (exception_id, exceptions)
+
+
+def exceptions_from_error_tuple(
+    exc_info,  # type: ExcInfo
+    client_options=None,  # type: Optional[Dict[str, Any]]
+    mechanism=None,  # type: Optional[Dict[str, Any]]
+):
+    # type: (...) -> List[Dict[str, Any]]
+    exc_type, exc_value, tb = exc_info
+
+    is_exception_group = BaseExceptionGroup is not None and isinstance(
+        exc_value, BaseExceptionGroup
+    )
+
+    if is_exception_group:
+        (_, exceptions) = exceptions_from_error(
+            exc_type=exc_type,
+            exc_value=exc_value,
+            tb=tb,
+            client_options=client_options,
+            mechanism=mechanism,
+            exception_id=0,
+            parent_id=0,
+        )
+
+    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.reverse()
+
+    return exceptions
+
+
+def to_string(value):
+    # type: (str) -> str
+    try:
+        return str(value)
+    except UnicodeDecodeError:
+        return repr(value)[1:-1]
+
+
+def iter_event_stacktraces(event):
+    # type: (Event) -> Iterator[Dict[str, Any]]
+    if "stacktrace" in event:
+        yield event["stacktrace"]
+    if "threads" in event:
+        for thread in event["threads"].get("values") or ():
+            if "stacktrace" in thread:
+                yield thread["stacktrace"]
+    if "exception" in event:
+        for exception in event["exception"].get("values") or ():
+            if "stacktrace" in exception:
+                yield exception["stacktrace"]
+
+
+def iter_event_frames(event):
+    # type: (Event) -> Iterator[Dict[str, Any]]
+    for stacktrace in iter_event_stacktraces(event):
+        for frame in stacktrace.get("frames") or ():
+            yield frame
+
+
+def handle_in_app(event, in_app_exclude=None, in_app_include=None, project_root=None):
+    # type: (Event, Optional[List[str]], Optional[List[str]], Optional[str]) -> Event
+    for stacktrace in iter_event_stacktraces(event):
+        set_in_app_in_frames(
+            stacktrace.get("frames"),
+            in_app_exclude=in_app_exclude,
+            in_app_include=in_app_include,
+            project_root=project_root,
+        )
+
+    return event
+
+
+def set_in_app_in_frames(frames, in_app_exclude, in_app_include, project_root=None):
+    # type: (Any, Optional[List[str]], Optional[List[str]], Optional[str]) -> Optional[Any]
+    if not frames:
+        return None
+
+    for frame in frames:
+        # if frame has already been marked as in_app, skip it
+        current_in_app = frame.get("in_app")
+        if current_in_app is not None:
+            continue
+
+        module = frame.get("module")
+
+        # check if module in frame is in the list of modules to include
+        if _module_in_list(module, in_app_include):
+            frame["in_app"] = True
+            continue
+
+        # check if module in frame is in the list of modules to exclude
+        if _module_in_list(module, in_app_exclude):
+            frame["in_app"] = False
+            continue
+
+        # if frame has no abs_path, skip further checks
+        abs_path = frame.get("abs_path")
+        if abs_path is None:
+            continue
+
+        if _is_external_source(abs_path):
+            frame["in_app"] = False
+            continue
+
+        if _is_in_project_root(abs_path, project_root):
+            frame["in_app"] = True
+            continue
+
+    return frames
+
+
+def exc_info_from_error(error):
+    # type: (Union[BaseException, ExcInfo]) -> ExcInfo
+    if isinstance(error, tuple) and len(error) == 3:
+        exc_type, exc_value, tb = error
+    elif isinstance(error, BaseException):
+        tb = getattr(error, "__traceback__", None)
+        if tb is not None:
+            exc_type = type(error)
+            exc_value = error
+        else:
+            exc_type, exc_value, tb = sys.exc_info()
+            if exc_value is not error:
+                tb = None
+                exc_value = error
+                exc_type = type(error)
+
+    else:
+        raise ValueError("Expected Exception object to report, got %s!" % type(error))
+
+    exc_info = (exc_type, exc_value, tb)
+
+    if TYPE_CHECKING:
+        # This cast is safe because exc_type and exc_value are either both
+        # None or both not None.
+        exc_info = cast(ExcInfo, exc_info)
+
+    return exc_info
+
+
+def event_from_exception(
+    exc_info,  # type: Union[BaseException, ExcInfo]
+    client_options=None,  # type: Optional[Dict[str, Any]]
+    mechanism=None,  # type: Optional[Dict[str, Any]]
+):
+    # type: (...) -> Tuple[Event, Dict[str, Any]]
+    exc_info = exc_info_from_error(exc_info)
+    hint = event_hint_with_exc_info(exc_info)
+    return (
+        {
+            "level": "error",
+            "exception": {
+                "values": exceptions_from_error_tuple(
+                    exc_info, client_options, mechanism
+                )
+            },
+        },
+        hint,
+    )
+
+
+def _module_in_list(name, items):
+    # type: (str | None, Optional[List[str]]) -> bool
+    if name is None:
+        return False
+
+    if not items:
+        return False
+
+    for item in items:
+        if item == name or name.startswith(item + "."):
+            return True
+
+    return False
+
+
+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
+    )
+    return external_source
+
+
+def _is_in_project_root(abs_path, project_root):
+    # type: (str, Optional[str]) -> bool
+    if project_root is None:
+        return False
+
+    # check if path is in the project root
+    if abs_path.startswith(project_root):
+        return True
+
+    return False
+
+
+def _truncate_by_bytes(string, max_bytes):
+    # type: (str, int) -> str
+    """
+    Truncate a UTF-8-encodable string to the last full codepoint so that it fits in max_bytes.
+    """
+    truncated = string.encode("utf-8")[: max_bytes - 3].decode("utf-8", errors="ignore")
+
+    return truncated + "..."
+
+
+def _get_size_in_bytes(value):
+    # type: (str) -> Optional[int]
+    try:
+        return len(value.encode("utf-8"))
+    except (UnicodeEncodeError, UnicodeDecodeError):
+        return None
+
+
+def strip_string(value, max_length=None):
+    # type: (str, Optional[int]) -> Union[AnnotatedValue, str]
+    if not value:
+        return value
+
+    if max_length is None:
+        max_length = DEFAULT_MAX_VALUE_LENGTH
+
+    byte_size = _get_size_in_bytes(value)
+    text_size = len(value)
+
+    if byte_size is not None and byte_size > max_length:
+        # truncate to max_length bytes, preserving code points
+        truncated_value = _truncate_by_bytes(value, max_length)
+    elif text_size is not None and text_size > max_length:
+        # fallback to truncating by string length
+        truncated_value = value[: max_length - 3] + "..."
+    else:
+        return value
+
+    return AnnotatedValue(
+        value=truncated_value,
+        metadata={
+            "len": byte_size or text_size,
+            "rem": [["!limit", "x", max_length - 3, max_length]],
+        },
+    )

posthog/feature_flags.py

@@ -7,12 +7,16 @@ 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)
 
 log = logging.getLogger("posthog")
 
+NONE_VALUES_ALLOWED_OPERATORS = ["is_not"]
+
 
 class InconclusiveMatchError(Exception):
     pass
@@ -22,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__
@@ -39,18 +43,29 @@ 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 []
+    valid_variant_keys = [variant["key"] for variant in flag_variants]
 
     # Stable sort conditions with variant overrides to the top. This ensures that if overrides are present, they are
     # evaluated first, and the variant override is applied to the first matching condition.
@@ -63,11 +78,11 @@ 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")
-                # Some filters can be explicitly set to null, which require accessing variants like so
-                flag_variants = ((flag.get("filters") or {}).get("multivariate") or {}).get("variants") or []
-                if variant_override and variant_override in [variant["key"] for variant in flag_variants]:
+                if variant_override and variant_override in valid_variant_keys:
                     variant = variant_override
                 else:
                     variant = get_matching_variant(flag, distinct_id)
@@ -76,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"):
@@ -98,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
@@ -112,19 +133,26 @@ 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")
 
     override_value = property_values[key]
 
+    if (operator not in NONE_VALUES_ALLOWED_OPERATORS) and override_value is None:
+        return False
+
     if operator in ("exact", "is_not"):
 
         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)
@@ -135,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,
@@ -175,27 +209,31 @@ def match_property(property, property_values) -> bool:
         else:
             return compare(str(override_value), str(value), operator)
 
-    if operator in ["is_date_before", "is_date_after", "is_relative_date_before", "is_relative_date_after"]:
+    if operator in ["is_date_before", "is_date_after"]:
         try:
-            if operator in ["is_relative_date_before", "is_relative_date_after"]:
-                parsed_date = relative_date_parse_for_feature_flag_matching(str(value))
-            else:
+            parsed_date = relative_date_parse_for_feature_flag_matching(str(value))
+
+            if not parsed_date:
                 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)
-            if operator in ("is_date_before", "is_relative_date_before"):
+            if operator == "is_date_before":
                 return override_date < parsed_date
             else:
                 return override_date > parsed_date
         elif isinstance(override_value, datetime.date):
-            if operator in ("is_date_before", "is_relative_date_before"):
+            if operator == "is_date_before":
                 return override_value < parsed_date.date()
             else:
                 return override_value > parsed_date.date()
@@ -203,14 +241,16 @@ def match_property(property, property_values) -> bool:
             try:
                 override_date = parser.parse(override_value)
                 override_date = convert_to_datetime_aware(override_date)
-                if operator in ("is_date_before", "is_relative_date_before"):
+                if operator == "is_date_before":
                     return override_date < parsed_date
                 else:
                     return override_date > parsed_date
             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}")
@@ -228,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)
@@ -264,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"
 
@@ -295,14 +339,18 @@ 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]:
-    regex = r"^(?P<number>[0-9]+)(?P<interval>[a-z])$"
+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)
     if match:

posthog/request.py

@@ -7,27 +7,57 @@ 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)
 
-DEFAULT_HOST = "https://app.posthog.com"
+US_INGESTION_ENDPOINT = "https://us.i.posthog.com"
+EU_INGESTION_ENDPOINT = "https://eu.i.posthog.com"
+DEFAULT_HOST = US_INGESTION_ENDPOINT
 USER_AGENT = "posthog-python/" + VERSION
 
 
+def determine_server_host(host: Optional[str]) -> str:
+    """Determines the server host to use."""
+    host_or_default = host or DEFAULT_HOST
+    trimmed_host = remove_trailing_slash(host_or_default)
+    if trimmed_host in ("https://app.posthog.com", "https://us.posthog.com"):
+        return US_INGESTION_ENDPOINT
+    elif trimmed_host == "https://eu.posthog.com":
+        return EU_INGESTION_ENDPOINT
+    else:
+        return host_or_default
+
+
 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")
     body = kwargs
-    body["sentAt"] = datetime.utcnow().replace(tzinfo=tzutc()).isoformat()
+    body["sentAt"] = datetime.now(tz=tzutc()).isoformat()
     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"
@@ -52,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)
@@ -61,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")
 
 
@@ -91,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,127 @@
+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, capture_exceptions=True):
+    """
+    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.
+        capture_exceptions: Whether to capture exceptions raised within the context (default: True).
+               If True, captures exceptions and tags them with the context tags before propagating them.
+               If False, exceptions will propagate without being tagged or captured.
+
+    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:
+        if capture_exceptions:
+            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, capture_exceptions=True):
+    """
+    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)
+        capture_exceptions: Whether to capture and track exceptions with posthog error tracking (default: True)
+
+    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, capture_exceptions=capture_exceptions):
+                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']}"
+                            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

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

posthog/test/ai/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/__init__.py

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

posthog/test/ai/openai/test_openai.py

@@ -0,0 +1,776 @@
+import json
+import time
+from unittest.mock import patch
+
+import pytest
+
+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,
+        ParsedResponse,
+    )
+    from openai.types.responses.parsed_response import (
+        ParsedResponseOutputMessage,
+        ParsedResponseOutputText,
+    )
+
+    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
+def mock_client():
+    with patch("posthog.client.Client") as mock_client:
+        mock_client.privacy_mode = False
+        yield mock_client
+
+
+@pytest.fixture
+def mock_openai_response():
+    return ChatCompletion(
+        id="test",
+        model="gpt-4",
+        object="chat.completion",
+        created=int(time.time()),
+        choices=[
+            Choice(
+                finish_reason="stop",
+                index=0,
+                message=ChatCompletionMessage(
+                    content="Test response",
+                    role="assistant",
+                ),
+            )
+        ],
+        usage=CompletionUsage(
+            completion_tokens=10,
+            prompt_tokens=20,
+            total_tokens=30,
+        ),
+    )
+
+
+@pytest.fixture
+def mock_openai_response_with_responses_api():
+    return Response(
+        id="test",
+        model="gpt-4o-mini",
+        object="response",
+        created_at=1741476542,
+        status="completed",
+        error=None,
+        incomplete_details=None,
+        instructions=None,
+        max_output_tokens=None,
+        tools=[],
+        tool_choice="auto",
+        output=[
+            ResponseOutputMessage(
+                id="msg_123",
+                type="message",
+                role="assistant",
+                status="completed",
+                content=[
+                    ResponseOutputText(
+                        type="output_text",
+                        text="Test response",
+                        annotations=[],
+                    )
+                ],
+            )
+        ],
+        parallel_tool_calls=True,
+        previous_response_id=None,
+        usage=ResponseUsage(
+            input_tokens=10,
+            output_tokens=10,
+            input_tokens_details={"prompt_tokens": 10, "cached_tokens": 0},
+            output_tokens_details={"reasoning_tokens": 15},
+            total_tokens=20,
+        ),
+        user=None,
+        metadata={},
+    )
+
+
+@pytest.fixture
+def mock_parsed_response():
+    return ParsedResponse(
+        id="test",
+        model="gpt-4o-2024-08-06",
+        object="response",
+        created_at=1741476542,
+        status="completed",
+        error=None,
+        incomplete_details=None,
+        instructions=None,
+        max_output_tokens=None,
+        tools=[],
+        tool_choice="auto",
+        output=[
+            ParsedResponseOutputMessage(
+                id="msg_123",
+                type="message",
+                role="assistant",
+                status="completed",
+                content=[
+                    ParsedResponseOutputText(
+                        type="output_text",
+                        text='{"name": "Science Fair", "date": "Friday", "participants": ["Alice", "Bob"]}',
+                        annotations=[],
+                        parsed={
+                            "name": "Science Fair",
+                            "date": "Friday",
+                            "participants": ["Alice", "Bob"],
+                        },
+                    )
+                ],
+            )
+        ],
+        output_parsed={
+            "name": "Science Fair",
+            "date": "Friday",
+            "participants": ["Alice", "Bob"],
+        },
+        parallel_tool_calls=True,
+        previous_response_id=None,
+        usage=ResponseUsage(
+            input_tokens=15,
+            output_tokens=20,
+            input_tokens_details={"prompt_tokens": 15, "cached_tokens": 0},
+            output_tokens_details={"reasoning_tokens": 5},
+            total_tokens=35,
+        ),
+        user=None,
+        metadata={},
+    )
+
+
+@pytest.fixture
+def mock_embedding_response():
+    return CreateEmbeddingResponse(
+        data=[
+            Embedding(
+                embedding=[0.1, 0.2, 0.3],
+                index=0,
+                object="embedding",
+            )
+        ],
+        model="text-embedding-3-small",
+        object="list",
+        usage=Usage(
+            prompt_tokens=10,
+            total_tokens=10,
+        ),
+    )
+
+
+@pytest.fixture
+def mock_openai_response_with_cached_tokens():
+    return ChatCompletion(
+        id="test",
+        model="gpt-4",
+        object="chat.completion",
+        created=int(time.time()),
+        choices=[
+            Choice(
+                finish_reason="stop",
+                index=0,
+                message=ChatCompletionMessage(
+                    content="Test response",
+                    role="assistant",
+                ),
+            )
+        ],
+        usage=CompletionUsage(
+            completion_tokens=10,
+            prompt_tokens=20,
+            total_tokens=30,
+            prompt_tokens_details={"cached_tokens": 15},
+        ),
+    )
+
+
+@pytest.fixture
+def mock_openai_response_with_tool_calls():
+    return ChatCompletion(
+        id="test",
+        model="gpt-4",
+        object="chat.completion",
+        created=int(time.time()),
+        choices=[
+            Choice(
+                finish_reason="tool_calls",
+                index=0,
+                message=ChatCompletionMessage(
+                    content="I'll check the weather for you.",
+                    role="assistant",
+                    tool_calls=[
+                        ChatCompletionMessageToolCall(
+                            id="call_abc123",
+                            type="function",
+                            function=Function(
+                                name="get_weather",
+                                arguments='{"location": "San Francisco", "unit": "celsius"}',
+                            ),
+                        )
+                    ],
+                ),
+            )
+        ],
+        usage=CompletionUsage(
+            completion_tokens=15,
+            prompt_tokens=20,
+            total_tokens=35,
+        ),
+    )
+
+
+def test_basic_completion(mock_client, mock_openai_response):
+    with patch(
+        "openai.resources.chat.completions.Completions.create",
+        return_value=mock_openai_response,
+    ):
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+        response = client.chat.completions.create(
+            model="gpt-4",
+            messages=[{"role": "user", "content": "Hello"}],
+            posthog_distinct_id="test-id",
+            posthog_properties={"foo": "bar"},
+        )
+
+        assert response == mock_openai_response
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+
+        assert call_args["distinct_id"] == "test-id"
+        assert call_args["event"] == "$ai_generation"
+        assert props["$ai_provider"] == "openai"
+        assert props["$ai_model"] == "gpt-4"
+        assert props["$ai_input"] == [{"role": "user", "content": "Hello"}]
+        assert props["$ai_output_choices"] == [
+            {"role": "assistant", "content": "Test response"}
+        ]
+        assert props["$ai_input_tokens"] == 20
+        assert props["$ai_output_tokens"] == 10
+        assert props["$ai_http_status"] == 200
+        assert props["foo"] == "bar"
+        assert isinstance(props["$ai_latency"], float)
+
+
+def test_embeddings(mock_client, mock_embedding_response):
+    with patch(
+        "openai.resources.embeddings.Embeddings.create",
+        return_value=mock_embedding_response,
+    ):
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+        response = client.embeddings.create(
+            model="text-embedding-3-small",
+            input="Hello world",
+            posthog_distinct_id="test-id",
+            posthog_properties={"foo": "bar"},
+        )
+
+        assert response == mock_embedding_response
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+
+        assert call_args["distinct_id"] == "test-id"
+        assert call_args["event"] == "$ai_embedding"
+        assert props["$ai_provider"] == "openai"
+        assert props["$ai_model"] == "text-embedding-3-small"
+        assert props["$ai_input"] == "Hello world"
+        assert props["$ai_input_tokens"] == 10
+        assert props["$ai_http_status"] == 200
+        assert props["foo"] == "bar"
+        assert isinstance(props["$ai_latency"], float)
+
+
+def test_groups(mock_client, mock_openai_response):
+    with patch(
+        "openai.resources.chat.completions.Completions.create",
+        return_value=mock_openai_response,
+    ):
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+        response = client.chat.completions.create(
+            model="gpt-4",
+            messages=[{"role": "user", "content": "Hello"}],
+            posthog_distinct_id="test-id",
+            posthog_groups={"company": "test_company"},
+        )
+
+        assert response == mock_openai_response
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+
+        assert call_args["groups"] == {"company": "test_company"}
+
+
+def test_privacy_mode_local(mock_client, mock_openai_response):
+    with patch(
+        "openai.resources.chat.completions.Completions.create",
+        return_value=mock_openai_response,
+    ):
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+        response = client.chat.completions.create(
+            model="gpt-4",
+            messages=[{"role": "user", "content": "Hello"}],
+            posthog_distinct_id="test-id",
+            posthog_privacy_mode=True,
+        )
+
+        assert response == mock_openai_response
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+        assert props["$ai_input"] is None
+        assert props["$ai_output_choices"] is None
+
+
+def test_privacy_mode_global(mock_client, mock_openai_response):
+    with patch(
+        "openai.resources.chat.completions.Completions.create",
+        return_value=mock_openai_response,
+    ):
+        mock_client.privacy_mode = True
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+        response = client.chat.completions.create(
+            model="gpt-4",
+            messages=[{"role": "user", "content": "Hello"}],
+            posthog_distinct_id="test-id",
+            posthog_privacy_mode=False,
+        )
+
+        assert response == mock_openai_response
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+        assert props["$ai_input"] is None
+        assert props["$ai_output_choices"] is None
+
+
+def test_error(mock_client, mock_openai_response):
+    with patch(
+        "openai.resources.chat.completions.Completions.create",
+        side_effect=Exception("Test error"),
+    ):
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+        with pytest.raises(Exception):
+            client.chat.completions.create(
+                model="gpt-4", messages=[{"role": "user", "content": "Hello"}]
+            )
+
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+        assert props["$ai_is_error"] is True
+        assert props["$ai_error"] == "Test error"
+
+
+def test_cached_tokens(mock_client, mock_openai_response_with_cached_tokens):
+    with patch(
+        "openai.resources.chat.completions.Completions.create",
+        return_value=mock_openai_response_with_cached_tokens,
+    ):
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+        response = client.chat.completions.create(
+            model="gpt-4",
+            messages=[{"role": "user", "content": "Hello"}],
+            posthog_distinct_id="test-id",
+            posthog_properties={"foo": "bar"},
+        )
+
+        assert response == mock_openai_response_with_cached_tokens
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+
+        assert call_args["distinct_id"] == "test-id"
+        assert call_args["event"] == "$ai_generation"
+        assert props["$ai_provider"] == "openai"
+        assert props["$ai_model"] == "gpt-4"
+        assert props["$ai_input"] == [{"role": "user", "content": "Hello"}]
+        assert props["$ai_output_choices"] == [
+            {"role": "assistant", "content": "Test response"}
+        ]
+        assert props["$ai_input_tokens"] == 20
+        assert props["$ai_output_tokens"] == 10
+        assert props["$ai_cache_read_input_tokens"] == 15
+        assert props["$ai_http_status"] == 200
+        assert props["foo"] == "bar"
+        assert isinstance(props["$ai_latency"], float)
+
+
+def test_tool_calls(mock_client, mock_openai_response_with_tool_calls):
+    with patch(
+        "openai.resources.chat.completions.Completions.create",
+        return_value=mock_openai_response_with_tool_calls,
+    ):
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+        response = client.chat.completions.create(
+            model="gpt-4",
+            messages=[
+                {"role": "user", "content": "What's the weather in San Francisco?"}
+            ],
+            tools=[
+                {
+                    "type": "function",
+                    "function": {
+                        "name": "get_weather",
+                        "description": "Get weather",
+                        "parameters": {},
+                    },
+                }
+            ],
+            posthog_distinct_id="test-id",
+        )
+
+        assert response == mock_openai_response_with_tool_calls
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+
+        assert call_args["distinct_id"] == "test-id"
+        assert call_args["event"] == "$ai_generation"
+        assert props["$ai_provider"] == "openai"
+        assert props["$ai_model"] == "gpt-4"
+        assert props["$ai_input"] == [
+            {"role": "user", "content": "What's the weather in San Francisco?"}
+        ]
+        assert props["$ai_output_choices"] == [
+            {"role": "assistant", "content": "I'll check the weather for you."}
+        ]
+
+        # Check that tool calls are properly captured
+        assert "$ai_tools" in props
+        tool_calls = props["$ai_tools"]
+        assert len(tool_calls) == 1
+
+        # Verify the tool call details
+        tool_call = tool_calls[0]
+        assert tool_call.id == "call_abc123"
+        assert tool_call.type == "function"
+        assert tool_call.function.name == "get_weather"
+
+        # Verify the arguments
+        arguments = tool_call.function.arguments
+        parsed_args = json.loads(arguments)
+        assert parsed_args == {"location": "San Francisco", "unit": "celsius"}
+
+        # Check token usage
+        assert props["$ai_input_tokens"] == 20
+        assert props["$ai_output_tokens"] == 15
+        assert props["$ai_http_status"] == 200
+
+
+def test_streaming_with_tool_calls(mock_client):
+    # Create mock tool call chunks that will be returned in sequence
+    tool_call_chunks = [
+        ChatCompletionChunk(
+            id="chunk1",
+            model="gpt-4",
+            object="chat.completion.chunk",
+            created=1234567890,
+            choices=[
+                ChoiceChunk(
+                    index=0,
+                    delta=ChoiceDelta(
+                        role="assistant",
+                        tool_calls=[
+                            ChoiceDeltaToolCall(
+                                index=0,
+                                id="call_abc123",
+                                type="function",
+                                function=ChoiceDeltaToolCallFunction(
+                                    name="get_weather",
+                                    arguments='{"location": "',
+                                ),
+                            )
+                        ],
+                    ),
+                    finish_reason=None,
+                )
+            ],
+        ),
+        ChatCompletionChunk(
+            id="chunk2",
+            model="gpt-4",
+            object="chat.completion.chunk",
+            created=1234567891,
+            choices=[
+                ChoiceChunk(
+                    index=0,
+                    delta=ChoiceDelta(
+                        tool_calls=[
+                            ChoiceDeltaToolCall(
+                                index=0,
+                                id="call_abc123",
+                                type="function",
+                                function=ChoiceDeltaToolCallFunction(
+                                    arguments='San Francisco"',
+                                ),
+                            )
+                        ],
+                    ),
+                    finish_reason=None,
+                )
+            ],
+        ),
+        ChatCompletionChunk(
+            id="chunk3",
+            model="gpt-4",
+            object="chat.completion.chunk",
+            created=1234567892,
+            choices=[
+                ChoiceChunk(
+                    index=0,
+                    delta=ChoiceDelta(
+                        tool_calls=[
+                            ChoiceDeltaToolCall(
+                                index=0,
+                                id="call_abc123",
+                                type="function",
+                                function=ChoiceDeltaToolCallFunction(
+                                    arguments=', "unit": "celsius"}',
+                                ),
+                            )
+                        ],
+                    ),
+                    finish_reason=None,
+                )
+            ],
+        ),
+        ChatCompletionChunk(
+            id="chunk4",
+            model="gpt-4",
+            object="chat.completion.chunk",
+            created=1234567893,
+            choices=[
+                ChoiceChunk(
+                    index=0,
+                    delta=ChoiceDelta(
+                        content="The weather in San Francisco is 15°C.",
+                    ),
+                    finish_reason=None,
+                )
+            ],
+            usage=CompletionUsage(
+                prompt_tokens=20,
+                completion_tokens=15,
+                total_tokens=35,
+            ),
+        ),
+    ]
+
+    # Mock the create method to return our chunks
+    with patch("openai.resources.chat.completions.Completions.create") as mock_create:
+        # Set up the mock to return our chunks when iterated
+        mock_create.return_value = tool_call_chunks
+
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+
+        # Call the streaming method
+        response_generator = client.chat.completions.create(
+            model="gpt-4",
+            messages=[
+                {"role": "user", "content": "What's the weather in San Francisco?"}
+            ],
+            tools=[
+                {
+                    "type": "function",
+                    "function": {
+                        "name": "get_weather",
+                        "description": "Get weather",
+                        "parameters": {},
+                    },
+                }
+            ],
+            stream=True,
+            posthog_distinct_id="test-id",
+        )
+
+        # Consume the generator to trigger the event capture
+        chunks = list(response_generator)
+
+        # Verify the chunks were returned correctly
+        assert len(chunks) == 4
+        assert chunks == tool_call_chunks
+
+        # Verify the capture was called with the right arguments
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+
+        assert call_args["distinct_id"] == "test-id"
+        assert call_args["event"] == "$ai_generation"
+        assert props["$ai_provider"] == "openai"
+        assert props["$ai_model"] == "gpt-4"
+
+        # Check that the tool calls were properly accumulated
+        assert "$ai_tools" in props
+        tool_calls = props["$ai_tools"]
+        assert len(tool_calls) == 1
+
+        # Verify the complete tool call was properly assembled
+        tool_call = tool_calls[0]
+        assert tool_call.id == "call_abc123"
+        assert tool_call.type == "function"
+        assert tool_call.function.name == "get_weather"
+
+        # Verify the arguments were concatenated correctly
+        arguments = tool_call.function.arguments
+        parsed_args = json.loads(arguments)
+        assert parsed_args == {"location": "San Francisco", "unit": "celsius"}
+
+        # Check that the content was also accumulated
+        assert (
+            props["$ai_output_choices"][0]["content"]
+            == "The weather in San Francisco is 15°C."
+        )
+
+        # Check token usage
+        assert props["$ai_input_tokens"] == 20
+        assert props["$ai_output_tokens"] == 15
+
+
+# test responses api
+def test_responses_api(mock_client, mock_openai_response_with_responses_api):
+    with patch(
+        "openai.resources.responses.Responses.create",
+        return_value=mock_openai_response_with_responses_api,
+    ):
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+        response = client.responses.create(
+            model="gpt-4o-mini",
+            input="Hello",
+            posthog_distinct_id="test-id",
+            posthog_properties={"foo": "bar"},
+        )
+        assert response == mock_openai_response_with_responses_api
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+
+        assert call_args["distinct_id"] == "test-id"
+        assert call_args["event"] == "$ai_generation"
+        assert props["$ai_provider"] == "openai"
+        assert props["$ai_model"] == "gpt-4o-mini"
+        assert props["$ai_input"] == [{"role": "user", "content": "Hello"}]
+        assert props["$ai_output_choices"] == [
+            {"role": "assistant", "content": "Test response"}
+        ]
+        assert props["$ai_input_tokens"] == 10
+        assert props["$ai_output_tokens"] == 10
+        assert props["$ai_reasoning_tokens"] == 15
+        assert props["$ai_http_status"] == 200
+        assert props["foo"] == "bar"
+        assert isinstance(props["$ai_latency"], float)
+
+
+def test_responses_parse(mock_client, mock_parsed_response):
+    with patch(
+        "openai.resources.responses.Responses.parse",
+        return_value=mock_parsed_response,
+    ):
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+        response = client.responses.parse(
+            model="gpt-4o-2024-08-06",
+            input=[
+                {"role": "system", "content": "Extract the event information."},
+                {
+                    "role": "user",
+                    "content": "Alice and Bob are going to a science fair on Friday.",
+                },
+            ],
+            text={
+                "format": {
+                    "type": "json_schema",
+                    "json_schema": {
+                        "name": "event",
+                        "schema": {
+                            "type": "object",
+                            "properties": {
+                                "name": {"type": "string"},
+                                "date": {"type": "string"},
+                                "participants": {
+                                    "type": "array",
+                                    "items": {"type": "string"},
+                                },
+                            },
+                            "required": ["name", "date", "participants"],
+                        },
+                    },
+                }
+            },
+            posthog_distinct_id="test-id",
+            posthog_properties={"foo": "bar"},
+        )
+
+        assert response == mock_parsed_response
+        assert mock_client.capture.call_count == 1
+
+        call_args = mock_client.capture.call_args[1]
+        props = call_args["properties"]
+
+        assert call_args["distinct_id"] == "test-id"
+        assert call_args["event"] == "$ai_generation"
+        assert props["$ai_provider"] == "openai"
+        assert props["$ai_model"] == "gpt-4o-2024-08-06"
+        assert props["$ai_input"] == [
+            {"role": "system", "content": "Extract the event information."},
+            {
+                "role": "user",
+                "content": "Alice and Bob are going to a science fair on Friday.",
+            },
+        ]
+        assert props["$ai_output_choices"] == [
+            {
+                "role": "assistant",
+                "content": '{"name": "Science Fair", "date": "Friday", "participants": ["Alice", "Bob"]}',
+            }
+        ]
+        assert props["$ai_input_tokens"] == 15
+        assert props["$ai_output_tokens"] == 20
+        assert props["$ai_reasoning_tokens"] == 5
+        assert props["$ai_http_status"] == 200
+        assert props["foo"] == "bar"
+        assert isinstance(props["$ai_latency"], float)

posthog/test/exception_integrations/__init__.py

@@ -0,0 +1,3 @@
+import pytest
+
+pytest.importorskip("django")

posthog/test/exception_integrations/test_django.py

@@ -0,0 +1,121 @@
+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"
+
+# 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):
+    factory = RequestFactory(
+        headers={
+            "User-Agent": DEFAULT_USER_AGENT,
+            "Referrer": "http://example.com",
+            "X-Forwarded-For": "193.4.5.12",
+            **(override_headers or {}),
+        }
+    )
+
+    request = factory.get("/api/endpoint")
+    return request
+
+
+def test_request_extractor_with_no_trace():
+    request = mock_request_factory(None)
+    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",
+    }
+
+
+def test_request_extractor_with_trace():
+    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",
+    }
+
+
+def test_request_extractor_with_tracestate():
+    request = mock_request_factory(
+        {
+            "traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01",
+            "tracestate": "posthog-distinct-id=1234",
+        }
+    )
+    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": "1234",
+        "$request_path": "/api/endpoint",
+    }
+
+
+def test_request_extractor_with_complicated_tracestate():
+    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_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_exception_capture.py

@@ -0,0 +1,63 @@
+import subprocess
+import sys
+from textwrap import dedent
+
+import pytest
+
+
+def test_excepthook(tmpdir):
+    app = tmpdir.join("app.py")
+    app.write(
+        dedent(
+            """
+    from posthog import Posthog
+    posthog = Posthog('phc_x', host='https://eu.i.posthog.com', enable_exception_autocapture=True, debug=True, on_error=lambda e, batch: print('error handling batch: ', e, batch))
+
+    # frame_value = "LOL"
+
+    1/0
+    """
+        )
+    )
+
+    with pytest.raises(subprocess.CalledProcessError) as excinfo:
+        subprocess.check_output([sys.executable, str(app)], stderr=subprocess.STDOUT)
+
+    output = excinfo.value.output
+
+    assert b"ZeroDivisionError" in output
+    assert b"LOL" in output
+    assert b"DEBUG:posthog:data uploaded successfully" in output
+    assert (
+        b'"$exception_list": [{"mechanism": {"type": "generic", "handled": true}, "module": null, "type": "ZeroDivisionError", "value": "division by zero", "stacktrace": {"frames": [{"platform": "python", "filename": "app.py", "abs_path"'
+        in output
+    )
+
+
+def test_trying_to_use_django_integration(tmpdir):
+    app = tmpdir.join("app.py")
+    app.write(
+        dedent(
+            """
+    from posthog import Posthog, Integrations
+    posthog = Posthog('phc_x', host='https://eu.i.posthog.com', enable_exception_autocapture=True, exception_autocapture_integrations=[Integrations.Django], debug=True, on_error=lambda e, batch: print('error handling batch: ', e, batch))
+
+    # frame_value = "LOL"
+
+    1/0
+    """
+        )
+    )
+
+    with pytest.raises(subprocess.CalledProcessError) as excinfo:
+        subprocess.check_output([sys.executable, str(app)], stderr=subprocess.STDOUT)
+
+    output = excinfo.value.output
+
+    assert b"ZeroDivisionError" in output
+    assert b"LOL" in output
+    assert b"DEBUG:posthog:data uploaded successfully" in output
+    assert (
+        b'"$exception_list": [{"mechanism": {"type": "generic", "handled": true}, "module": null, "type": "ZeroDivisionError", "value": "division by zero", "stacktrace": {"frames": [{"platform": "python", "filename": "app.py", "abs_path"'
+        in output
+    )

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

@@ -6,12 +6,18 @@ from posthog import Posthog
 class TestModule(unittest.TestCase):
     posthog = None
 
+    def _assert_enqueue_result(self, result):
+        self.assertEqual(type(result[0]), bool)
+        self.assertEqual(type(result[1]), dict)
+
     def failed(self):
         self.failed = True
 
     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
@@ -22,15 +28,18 @@ class TestModule(unittest.TestCase):
         self.assertRaises(Exception, self.posthog.capture)
 
     def test_track(self):
-        self.posthog.capture("distinct_id", "python module event")
+        res = self.posthog.capture("distinct_id", "python module event")
+        self._assert_enqueue_result(res)
         self.posthog.flush()
 
     def test_identify(self):
-        self.posthog.identify("distinct_id", {"email": "user@email.com"})
+        res = self.posthog.identify("distinct_id", {"email": "user@email.com"})
+        self._assert_enqueue_result(res)
         self.posthog.flush()
 
     def test_alias(self):
-        self.posthog.alias("previousId", "distinct_id")
+        res = self.posthog.alias("previousId", "distinct_id")
+        self._assert_enqueue_result(res)
         self.posthog.flush()
 
     def test_page(self):

posthog/test/test_request.py

@@ -2,22 +2,39 @@ import json
 import unittest
 from datetime import date, datetime
 
+import mock
+import pytest
 import requests
 
-from posthog.request import DatetimeSerializer, batch_post
+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)}
@@ -33,12 +50,81 @@ 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",
+    [
+        ("https://t.posthog.com", "https://t.posthog.com"),
+        ("https://t.posthog.com/", "https://t.posthog.com/"),
+        ("t.posthog.com", "t.posthog.com"),
+        ("t.posthog.com/", "t.posthog.com/"),
+        ("https://us.posthog.com.rg.proxy.com", "https://us.posthog.com.rg.proxy.com"),
+        ("app.posthog.com", "app.posthog.com"),
+        ("eu.posthog.com", "eu.posthog.com"),
+        ("https://app.posthog.com", "https://us.i.posthog.com"),
+        ("https://eu.posthog.com", "https://eu.i.posthog.com"),
+        ("https://us.posthog.com", "https://us.i.posthog.com"),
+        ("https://app.posthog.com/", "https://us.i.posthog.com"),
+        ("https://eu.posthog.com/", "https://eu.i.posthog.com"),
+        ("https://us.posthog.com/", "https://us.i.posthog.com"),
+        (None, "https://us.i.posthog.com"),
+    ],
+)
+def test_routing_to_custom_host(host, expected):
+    assert determine_server_host(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,10 +1,15 @@
 import unittest
+from dataclasses import dataclass
 from datetime import date, datetime, timedelta
 from decimal import Decimal
+from typing import Optional
 from uuid import UUID
 
 import six
 from dateutil.tz import tzutc
+from parameterized import parameterized
+from pydantic import BaseModel
+from pydantic.v1 import BaseModel as BaseModelV1
 
 from posthog import utils
 
@@ -13,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 = {
@@ -50,52 +67,109 @@ 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):
+            foo: str
+            bar: int
+            baz: Optional[str] = None
 
-class TestSizeLimitedDict(unittest.TestCase):
-    def test_size_limited_dict(self):
-        size = 10
-        values = utils.SizeLimitedDict(size, lambda _: -1)
+        class ModelV1(BaseModelV1):
+            foo: int
+            bar: str
 
-        for i in range(100):
-            values[i] = i
+        class NestedModel(BaseModel):
+            foo: ModelV2
 
-            self.assertEqual(values[i], i)
-            self.assertEqual(len(values), i % size + 1)
+        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"}
+        }
 
-            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))
+    def test_clean_pydantic_like_class(self) -> None:
+        class Dummy:
+            def model_dump(self, required_param: str) -> dict:
+                return {}
+
+        # 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
+
+        @dataclass
+        class TestDataClass:
+            foo: str
+            bar: int
+            nested: InnerDataClass
+
+        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,14 +53,28 @@ def clean(item):
         return float(item)
     if isinstance(item, UUID):
         return str(item)
-    elif isinstance(item, (six.string_types, bool, numbers.Number, datetime, date, type(None))):
+    if isinstance(
+        item, (six.string_types, bool, numbers.Number, datetime, date, type(None))
+    ):
         return item
-    elif isinstance(item, (set, list, tuple)):
+    if isinstance(item, (set, list, tuple)):
         return _clean_list(item)
-    elif isinstance(item, dict):
+    # Pydantic model
+    try:
+        # v2+
+        if hasattr(item, "model_dump") and callable(item.model_dump):
+            item = item.model_dump()
+        # v1
+        elif hasattr(item, "dict") and callable(item.dict):
+            item = item.dict()
+    except TypeError as e:
+        log.debug(f"Could not serialize Pydantic-like model: {e}")
+        pass
+    if isinstance(item, dict):
         return _clean_dict(item)
-    else:
-        return _coerce_unicode(item)
+    if is_dataclass(item) and not isinstance(item, type):
+        return _clean_dataclass(item)
+    return _coerce_unicode(item)
 
 
 def _clean_list(list_):
@@ -80,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
 
 
@@ -115,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.2.0"
+VERSION = "4.8.0"
 
 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

sentry_django_example/sentry_django_example/urls.py

@@ -13,6 +13,7 @@ Including another URLconf
     1. Import the include() function: from django.urls import include, path
     2. Add a URL to urlpatterns:  path('blog/', include('blog.urls'))
 """
+
 from django.contrib import admin
 from django.urls import path
 

setup.cfg

@@ -1,2 +0,0 @@
-[bdist_wheel]
-universal = 1

setup.py

@@ -8,57 +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"],
-    "sentry": ["sentry-sdk", "django"],
-}
-
+# 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.test", "posthog.sentry"],
     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",
-    ],
+    # This will fallback to pyproject.toml for detailed configuration
 )

setup_analytics.py

@@ -1,57 +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.test", "posthoganalytics.sentry"],
     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",
-    ],
+    # This will fallback to pyproject.toml for detailed configuration
 )

simulator.py

@@ -24,12 +24,13 @@ parser.add_argument("--type", help="The posthog message type")
 
 parser.add_argument("--distinct_id", help="the user id to send the event as")
 parser.add_argument("--anonymousId", help="the anonymous user id to send the event as")
-parser.add_argument("--context", help="additional context for the event (JSON-encoded)")
 
 parser.add_argument("--event", help="the event name to send with the event")
 parser.add_argument("--properties", help="the event properties to send (JSON-encoded)")
 
-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)")
 
@@ -48,7 +49,6 @@ def capture():
         options.event,
         anonymous_id=options.anonymousId,
         properties=json_hash(options.properties),
-        context=json_hash(options.context),
     )
 
 
@@ -58,7 +58,6 @@ def page():
         name=options.name,
         anonymous_id=options.anonymousId,
         properties=json_hash(options.properties),
-        context=json_hash(options.context),
     )
 
 
@@ -67,7 +66,6 @@ def identify():
         options.distinct_id,
         anonymous_id=options.anonymousId,
         traits=json_hash(options.traits),
-        context=json_hash(options.context),
     )
 
 
@@ -75,7 +73,6 @@ def set_once():
     posthog.set_once(
         options.distinct_id,
         properties=json_hash(options.traits),
-        context=json_hash(options.context),
     )
 
 
@@ -83,7 +80,6 @@ def set():
     posthog.set(
         options.distinct_id,
         properties=json_hash(options.traits),
-        context=json_hash(options.context),
     )
 
 
@@ -100,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: