feat: azure export w/ async (#200)

* feat: azure export w/ async

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

.github/workflows/ci.yml

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

.gitignore

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

CHANGELOG.md

@@ -1,3 +1,143 @@
+
+## 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.
@@ -38,7 +178,6 @@
 
 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.
@@ -64,6 +203,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:
@@ -104,7 +244,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
@@ -120,6 +259,7 @@ Changes:
 Changes:
 
 1. Fixes issues with date comparison.
+
 ## 2.1.1 - 2022-09-14
 
 Changes:
@@ -132,6 +272,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
@@ -144,7 +285,7 @@ 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**
+   **Note: These are removed in 2.0.2**
 3. Feature flag remote evaluation doesn't require a personal API key.
 
 New Changes:
@@ -152,18 +293,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

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

Makefile

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

README.md

@@ -10,8 +10,10 @@ Please see the [Python integration docs](https://posthog.com/docs/integrations/p
 ### Testing Locally
 
 1. Run `python3 -m venv env` (creates virtual environment called "env")
+    * or `uv venv env`
 2. Run `source env/bin/activate` (activates the virtual environment)
 3. Run `python3 -m pip install -e ".[test]"` (installs the package in develop mode, along with test dependencies)
+    * or `uv pip install -e ".[test]"`
 4. Run `make test`
   1. To run a specific test do `pytest -k test_no_api_key`
 

example.py

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

llm_observability_examples.py

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

mypy-baseline.txt

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

mypy.ini

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

posthog/__init__.py

@@ -1,7 +1,9 @@
 import datetime  # noqa: F401
-from typing import Callable, Dict, Optional, Tuple  # 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.version import VERSION
 
 __version__ = VERSION
@@ -19,8 +21,14 @@ 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]
+# Used to determine in app paths for exception autocapture. Defaults to the current working directory
+project_root = None  # type: Optional[str]
+# Used for our AI observability feature to not capture any prompt or output just usage + metadata
+privacy_mode = False  # type: bool
 
 default_client = None  # type: Optional[Client]
 
@@ -57,6 +65,14 @@ def capture(
     posthog.capture('distinct id', 'purchase', groups={'company': 'id:5'})
     ```
     """
+
+    if context is not None:
+        warnings.warn(
+            "The 'context' parameter is deprecated and will be removed in a future version.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
     return _proxy(
         "capture",
         distinct_id=distinct_id,
@@ -95,6 +111,14 @@ def identify(
     })
     ```
     """
+
+    if context is not None:
+        warnings.warn(
+            "The 'context' parameter is deprecated and will be removed in a future version.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
     return _proxy(
         "identify",
         distinct_id=distinct_id,
@@ -130,6 +154,14 @@ def set(
      })
      ```
     """
+
+    if context is not None:
+        warnings.warn(
+            "The 'context' parameter is deprecated and will be removed in a future version.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
     return _proxy(
         "set",
         distinct_id=distinct_id,
@@ -165,6 +197,14 @@ def set_once(
      })
      ```
     """
+
+    if context is not None:
+        warnings.warn(
+            "The 'context' parameter is deprecated and will be removed in a future version.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
     return _proxy(
         "set_once",
         distinct_id=distinct_id,
@@ -201,6 +241,14 @@ def group_identify(
      })
      ```
     """
+
+    if context is not None:
+        warnings.warn(
+            "The 'context' parameter is deprecated and will be removed in a future version.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
     return _proxy(
         "group_identify",
         group_type=group_type,
@@ -238,6 +286,14 @@ def alias(
     posthog.alias('anonymous session id', 'distinct id')
     ```
     """
+
+    if context is not None:
+        warnings.warn(
+            "The 'context' parameter is deprecated and will be removed in a future version.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
     return _proxy(
         "alias",
         previous_id=previous_id,
@@ -249,6 +305,58 @@ 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]
+):
+    # 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
+
+    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,
+    )
+
+
 def feature_enabled(
     key,  # type: str
     distinct_id,  # type: str
@@ -384,6 +492,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={},
@@ -456,10 +584,12 @@ def _proxy(method, *args, **kwargs):
             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,
+            exception_autocapture_integrations=exception_autocapture_integrations,
         )
 
     # always set incase user changes it

posthog/ai/anthropic/__init__.py

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

posthog/ai/anthropic/anthropic.py

@@ -0,0 +1,202 @@
+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 = 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 = uuid.uuid4()
+
+        return self._create_streaming(
+            posthog_distinct_id,
+            posthog_trace_id,
+            posthog_properties,
+            posthog_privacy_mode,
+            posthog_groups,
+            **kwargs,
+        )
+
+    def _create_streaming(
+        self,
+        posthog_distinct_id: Optional[str],
+        posthog_trace_id: Optional[str],
+        posthog_properties: Optional[Dict[str, Any]],
+        posthog_privacy_mode: bool,
+        posthog_groups: Optional[Dict[str, Any]],
+        **kwargs: Any,
+    ):
+        start_time = time.time()
+        usage_stats: Dict[str, int] = {"input_tokens": 0, "output_tokens": 0}
+        accumulated_content = []
+        response = super().create(**kwargs)
+
+        def generator():
+            nonlocal usage_stats
+            nonlocal accumulated_content
+            try:
+                for event in response:
+                    if hasattr(event, "usage") and event.usage:
+                        usage_stats = {
+                            k: getattr(event.usage, k, 0)
+                            for k in [
+                                "input_tokens",
+                                "output_tokens",
+                            ]
+                        }
+
+                    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 = 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_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,202 @@
+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 = 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 = uuid.uuid4()
+
+        return await self._create_streaming(
+            posthog_distinct_id,
+            posthog_trace_id,
+            posthog_properties,
+            posthog_privacy_mode,
+            posthog_groups,
+            **kwargs,
+        )
+
+    async def _create_streaming(
+        self,
+        posthog_distinct_id: Optional[str],
+        posthog_trace_id: Optional[str],
+        posthog_properties: Optional[Dict[str, Any]],
+        posthog_privacy_mode: bool,
+        posthog_groups: Optional[Dict[str, Any]],
+        **kwargs: Any,
+    ):
+        start_time = time.time()
+        usage_stats: Dict[str, int] = {"input_tokens": 0, "output_tokens": 0}
+        accumulated_content = []
+        response = await super().create(**kwargs)
+
+        async def generator():
+            nonlocal usage_stats
+            nonlocal accumulated_content
+            try:
+                async for event in response:
+                    if hasattr(event, "usage") and event.usage:
+                        usage_stats = {
+                            k: getattr(event.usage, k, 0)
+                            for k in [
+                                "input_tokens",
+                                "output_tokens",
+                            ]
+                        }
+
+                    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 = 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_latency": latency,
+            "$ai_trace_id": posthog_trace_id,
+            "$ai_base_url": str(self._client.base_url),
+            **(posthog_properties or {}),
+        }
+
+        if posthog_distinct_id is None:
+            event_properties["$process_person_profile"] = False
+
+        if hasattr(self._client._ph_client, "capture"):
+            self._client._ph_client.capture(
+                distinct_id=posthog_distinct_id or posthog_trace_id,
+                event="$ai_generation",
+                properties=event_properties,
+                groups=posthog_groups,
+            )

posthog/ai/anthropic/anthropic_providers.py

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

posthog/ai/langchain/__init__.py

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

posthog/ai/langchain/callbacks.py

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

posthog/ai/openai/openai_async.py

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

posthog/ai/openai/openai_providers.py

@@ -0,0 +1,41 @@
+try:
+    import openai
+    import openai.resources
+except ImportError:
+    raise ModuleNotFoundError("Please install the Open AI SDK to use this feature: 'pip install openai'")
+
+from posthog.ai.openai.openai import WrappedBeta, WrappedChat, WrappedEmbeddings
+from posthog.ai.openai.openai_async import WrappedBeta as AsyncWrappedBeta
+from posthog.ai.openai.openai_async import WrappedChat as AsyncWrappedChat
+from posthog.ai.openai.openai_async import WrappedEmbeddings as AsyncWrappedEmbeddings
+from posthog.client import Client as PostHogClient
+
+
+class AzureOpenAI(openai.AzureOpenAI):
+    """
+    A wrapper around the Azure OpenAI SDK that automatically sends LLM usage events to PostHog.
+    """
+
+    _ph_client: PostHogClient
+
+    def __init__(self, posthog_client: PostHogClient, **kwargs):
+        super().__init__(**kwargs)
+        self._ph_client = posthog_client
+        self.chat = WrappedChat(self)
+        self.embeddings = WrappedEmbeddings(self)
+        self.beta = WrappedBeta(self)
+
+
+class AsyncAzureOpenAI(openai.AsyncAzureOpenAI):
+    """
+    A wrapper around the Azure OpenAI SDK that automatically sends LLM usage events to PostHog.
+    """
+
+    _ph_client: PostHogClient
+
+    def __init__(self, posthog_client: PostHogClient, **kwargs):
+        super().__init__(**kwargs)
+        self._ph_client = posthog_client
+        self.chat = AsyncWrappedChat(self)
+        self.embeddings = AsyncWrappedEmbeddings(self)
+        self.beta = AsyncWrappedBeta(self)

posthog/ai/utils.py

@@ -0,0 +1,257 @@
+import time
+import uuid
+from typing import Any, Callable, Dict, 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,
+        }
+    elif provider == "openai":
+        return {
+            "input_tokens": response.usage.prompt_tokens,
+            "output_tokens": response.usage.completion_tokens,
+        }
+    return {
+        "input_tokens": 0,
+        "output_tokens": 0,
+    }
+
+
+def format_response(response, provider: str):
+    """
+    Format a regular (non-streaming) response.
+    """
+    output = []
+    if response is None:
+        return output
+    if provider == "anthropic":
+        return format_response_anthropic(response)
+    elif provider == "openai":
+        return format_response_openai(response)
+    return output
+
+
+def format_response_anthropic(response):
+    output = []
+    for choice in response.content:
+        if choice.text:
+            output.append(
+                {
+                    "role": "assistant",
+                    "content": choice.text,
+                }
+            )
+    return output
+
+
+def format_response_openai(response):
+    output = []
+    for choice in response.choices:
+        if choice.message.content:
+            output.append(
+                {
+                    "content": choice.message.content,
+                    "role": choice.message.role,
+                }
+            )
+    return output
+
+
+def merge_system_prompt(kwargs: Dict[str, Any], provider: str):
+    if provider != "anthropic":
+        return kwargs.get("messages")
+    messages = kwargs.get("messages") or []
+    if kwargs.get("system") is None:
+        return messages
+    return [{"role": "system", "content": kwargs.get("system")}] + messages
+
+
+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 = uuid.uuid4()
+
+        if response and hasattr(response, "usage"):
+            usage = get_usage(response, provider)
+
+        messages = merge_system_prompt(kwargs, provider)
+
+        event_properties = {
+            "$ai_provider": provider,
+            "$ai_model": kwargs.get("model"),
+            "$ai_model_parameters": get_model_params(kwargs),
+            "$ai_input": with_privacy_mode(ph_client, posthog_privacy_mode, messages),
+            "$ai_output_choices": with_privacy_mode(
+                ph_client, posthog_privacy_mode, format_response(response, provider)
+            ),
+            "$ai_http_status": http_status,
+            "$ai_input_tokens": usage.get("input_tokens", 0),
+            "$ai_output_tokens": usage.get("output_tokens", 0),
+            "$ai_latency": latency,
+            "$ai_trace_id": posthog_trace_id,
+            "$ai_base_url": str(base_url),
+            **(posthog_properties or {}),
+            **(error_params or {}),
+        }
+
+        if posthog_distinct_id is None:
+            event_properties["$process_person_profile"] = False
+
+        # 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 = uuid.uuid4()
+
+        if response and hasattr(response, "usage"):
+            usage = get_usage(response, provider)
+
+        messages = merge_system_prompt(kwargs, provider)
+
+        event_properties = {
+            "$ai_provider": provider,
+            "$ai_model": kwargs.get("model"),
+            "$ai_model_parameters": get_model_params(kwargs),
+            "$ai_input": with_privacy_mode(ph_client, posthog_privacy_mode, messages),
+            "$ai_output_choices": with_privacy_mode(
+                ph_client, posthog_privacy_mode, format_response(response, provider)
+            ),
+            "$ai_http_status": http_status,
+            "$ai_input_tokens": usage.get("input_tokens", 0),
+            "$ai_output_tokens": usage.get("output_tokens", 0),
+            "$ai_latency": latency,
+            "$ai_trace_id": posthog_trace_id,
+            "$ai_base_url": str(base_url),
+            **(posthog_properties or {}),
+            **(error_params or {}),
+        }
+
+        if posthog_distinct_id is None:
+            event_properties["$process_person_profile"] = False
+
+        # send the event to posthog
+        if hasattr(ph_client, "capture") and callable(ph_client.capture):
+            ph_client.capture(
+                distinct_id=posthog_distinct_id or posthog_trace_id,
+                event="$ai_generation",
+                properties=event_properties,
+                groups=posthog_groups,
+            )
+
+    if error:
+        raise error
+
+    return response
+
+
+def with_privacy_mode(ph_client: PostHogClient, privacy_mode: bool, value: Any):
+    if ph_client.privacy_mode or privacy_mode:
+        return None
+    return value

posthog/client.py

@@ -1,18 +1,25 @@
 import atexit
 import logging
 import numbers
+import os
+import platform
+import sys
+import warnings
 from datetime import datetime, timedelta
-from uuid import UUID
+from typing import Any
+from uuid import UUID, uuid4
 
+import distro  # For Linux OS detection
 from dateutil.tz import tzutc
 from six import string_types
 
 from posthog.consumer import Consumer
 from posthog.exception_capture import ExceptionCapture
+from posthog.exception_utils import exc_info_from_error, exceptions_from_error_tuple, handle_in_app
 from posthog.feature_flags import InconclusiveMatchError, match_feature_flag_properties
 from posthog.poller import Poller
-from posthog.request import APIError, batch_post, decide, determine_server_host, get
-from posthog.utils import SizeLimitedDict, clean, guess_timezone
+from posthog.request import DEFAULT_HOST, APIError, batch_post, decide, determine_server_host, get, remote_config
+from posthog.utils import SizeLimitedDict, clean, guess_timezone, remove_trailing_slash
 from posthog.version import VERSION
 
 try:
@@ -25,6 +32,60 @@ ID_TYPES = (numbers.Number, string_types, UUID)
 MAX_DICT_SIZE = 50_000
 
 
+def get_os_info():
+    """
+    Returns standardized OS name and version information.
+    Similar to how user agent parsing works in JS.
+    """
+    os_name = ""
+    os_version = ""
+
+    platform_name = sys.platform
+
+    if platform_name.startswith("win"):
+        os_name = "Windows"
+        if hasattr(platform, "win32_ver"):
+            win_version = platform.win32_ver()[0]
+            if win_version:
+                os_version = win_version
+
+    elif platform_name == "darwin":
+        os_name = "Mac OS X"
+        if hasattr(platform, "mac_ver"):
+            mac_version = platform.mac_ver()[0]
+            if mac_version:
+                os_version = mac_version
+
+    elif platform_name.startswith("linux"):
+        os_name = "Linux"
+        linux_info = distro.info()
+        if linux_info["version"]:
+            os_version = linux_info["version"]
+
+    elif platform_name.startswith("freebsd"):
+        os_name = "FreeBSD"
+        if hasattr(platform, "release"):
+            os_version = platform.release()
+
+    else:
+        os_name = platform_name
+        if hasattr(platform, "release"):
+            os_version = platform.release()
+
+    return os_name, os_version
+
+
+def system_context() -> dict[str, Any]:
+    os_name, os_version = get_os_info()
+
+    return {
+        "$python_runtime": platform.python_implementation(),
+        "$python_version": "%s.%s.%s" % (sys.version_info[:3]),
+        "$os": os_name,
+        "$os_version": os_version,
+    }
+
+
 class Client(object):
     """Create a new PostHog client."""
 
@@ -52,7 +113,11 @@ class Client(object):
         disable_geoip=True,
         historical_migration=False,
         feature_flags_request_timeout_seconds=3,
+        super_properties=None,
         enable_exception_autocapture=False,
+        exception_autocapture_integrations=None,
+        project_root=None,
+        privacy_mode=False,
     ):
         self.queue = queue.Queue(max_queue_size)
 
@@ -65,6 +130,8 @@ class Client(object):
         self.debug = debug
         self.send = send
         self.sync_mode = sync_mode
+        # Used for session replay URL generation - we don't want the server host here.
+        self.raw_host = host or DEFAULT_HOST
         self.host = determine_server_host(host)
         self.gzip = gzip
         self.timeout = timeout
@@ -79,7 +146,19 @@ class Client(object):
         self.disabled = disabled
         self.disable_geoip = disable_geoip
         self.historical_migration = historical_migration
+        self.super_properties = super_properties
         self.enable_exception_autocapture = enable_exception_autocapture
+        self.exception_autocapture_integrations = exception_autocapture_integrations
+        self.exception_capture = None
+        self.privacy_mode = privacy_mode
+
+        if project_root is None:
+            try:
+                project_root = os.getcwd()
+            except Exception:
+                project_root = None
+
+        self.project_root = project_root
 
         # personal_api_key: This should be a generated Personal API Key, private
         self.personal_api_key = personal_api_key
@@ -92,7 +171,7 @@ class Client(object):
             self.log.setLevel(logging.WARNING)
 
         if self.enable_exception_autocapture:
-            self.exception_capture = ExceptionCapture(self)
+            self.exception_capture = ExceptionCapture(self, integrations=self.exception_autocapture_integrations)
 
         if sync_mode:
             self.consumers = None
@@ -126,14 +205,19 @@ class Client(object):
                     consumer.start()
 
     def identify(self, distinct_id=None, properties=None, context=None, timestamp=None, uuid=None, disable_geoip=None):
+        if context is not None:
+            warnings.warn(
+                "The 'context' parameter is deprecated and will be removed in a future version.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
         properties = properties or {}
-        context = context or {}
         require("distinct_id", distinct_id, ID_TYPES)
         require("properties", properties, dict)
 
         msg = {
             "timestamp": timestamp,
-            "context": context,
             "distinct_id": distinct_id,
             "$set": properties,
             "event": "$identify",
@@ -154,6 +238,15 @@ class Client(object):
         resp_data = self.get_decide(distinct_id, groups, person_properties, group_properties, disable_geoip)
         return resp_data["featureFlagPayloads"]
 
+    def get_feature_flags_and_payloads(
+        self, distinct_id, groups=None, person_properties=None, group_properties=None, disable_geoip=None
+    ):
+        resp_data = self.get_decide(distinct_id, groups, person_properties, group_properties, disable_geoip)
+        return {
+            "featureFlags": resp_data["featureFlags"],
+            "featureFlagPayloads": resp_data["featureFlagPayloads"],
+        }
+
     def get_decide(self, distinct_id, groups=None, person_properties=None, group_properties=None, disable_geoip=None):
         require("distinct_id", distinct_id, ID_TYPES)
 
@@ -188,8 +281,15 @@ class Client(object):
         send_feature_flags=False,
         disable_geoip=None,
     ):
-        properties = properties or {}
-        context = context or {}
+        if context is not None:
+            warnings.warn(
+                "The 'context' parameter is deprecated and will be removed in a future version.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
+        properties = {**(properties or {}), **system_context()}
+
         require("distinct_id", distinct_id, ID_TYPES)
         require("properties", properties, dict)
         require("event", event, string_types)
@@ -197,7 +297,6 @@ class Client(object):
         msg = {
             "properties": properties,
             "timestamp": timestamp,
-            "context": context,
             "distinct_id": distinct_id,
             "event": event,
             "uuid": uuid,
@@ -215,7 +314,7 @@ class Client(object):
             except Exception as e:
                 self.log.exception(f"[FEATURE FLAGS] Unable to get feature variants: {e}")
 
-        elif self.feature_flags:
+        elif self.feature_flags and event != "$feature_flag_called":
             # Local evaluation is enabled, flags are loaded, so try and get all flags we can without going to the server
             feature_variants = self.get_all_flags(
                 distinct_id, groups=(groups or {}), disable_geoip=disable_geoip, only_evaluate_locally=True
@@ -234,14 +333,19 @@ class Client(object):
         return self._enqueue(msg, disable_geoip)
 
     def set(self, distinct_id=None, properties=None, context=None, timestamp=None, uuid=None, disable_geoip=None):
+        if context is not None:
+            warnings.warn(
+                "The 'context' parameter is deprecated and will be removed in a future version.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
         properties = properties or {}
-        context = context or {}
         require("distinct_id", distinct_id, ID_TYPES)
         require("properties", properties, dict)
 
         msg = {
             "timestamp": timestamp,
-            "context": context,
             "distinct_id": distinct_id,
             "$set": properties,
             "event": "$set",
@@ -251,14 +355,19 @@ class Client(object):
         return self._enqueue(msg, disable_geoip)
 
     def set_once(self, distinct_id=None, properties=None, context=None, timestamp=None, uuid=None, disable_geoip=None):
+        if context is not None:
+            warnings.warn(
+                "The 'context' parameter is deprecated and will be removed in a future version.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
         properties = properties or {}
-        context = context or {}
         require("distinct_id", distinct_id, ID_TYPES)
         require("properties", properties, dict)
 
         msg = {
             "timestamp": timestamp,
-            "context": context,
             "distinct_id": distinct_id,
             "$set_once": properties,
             "event": "$set_once",
@@ -276,13 +385,24 @@ class Client(object):
         timestamp=None,
         uuid=None,
         disable_geoip=None,
+        distinct_id=None,
     ):
+        if context is not None:
+            warnings.warn(
+                "The 'context' parameter is deprecated and will be removed in a future version.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
         properties = properties or {}
-        context = context or {}
         require("group_type", group_type, ID_TYPES)
         require("group_key", group_key, ID_TYPES)
         require("properties", properties, dict)
 
+        if distinct_id:
+            require("distinct_id", distinct_id, ID_TYPES)
+        else:
+            distinct_id = "${}_{}".format(group_type, group_key)
+
         msg = {
             "event": "$groupidentify",
             "properties": {
@@ -290,16 +410,20 @@ class Client(object):
                 "$group_key": group_key,
                 "$group_set": properties,
             },
-            "distinct_id": "${}_{}".format(group_type, group_key),
+            "distinct_id": distinct_id,
             "timestamp": timestamp,
-            "context": context,
             "uuid": uuid,
         }
 
         return self._enqueue(msg, disable_geoip)
 
     def alias(self, previous_id=None, distinct_id=None, context=None, timestamp=None, uuid=None, disable_geoip=None):
-        context = context or {}
+        if context is not None:
+            warnings.warn(
+                "The 'context' parameter is deprecated and will be removed in a future version.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
 
         require("previous_id", previous_id, ID_TYPES)
         require("distinct_id", distinct_id, ID_TYPES)
@@ -310,7 +434,6 @@ class Client(object):
                 "alias": distinct_id,
             },
             "timestamp": timestamp,
-            "context": context,
             "event": "$create_alias",
             "distinct_id": previous_id,
         }
@@ -320,9 +443,14 @@ class Client(object):
     def page(
         self, distinct_id=None, url=None, properties=None, context=None, timestamp=None, uuid=None, disable_geoip=None
     ):
-        properties = properties or {}
-        context = context or {}
+        if context is not None:
+            warnings.warn(
+                "The 'context' parameter is deprecated and will be removed in a future version.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
 
+        properties = properties or {}
         require("distinct_id", distinct_id, ID_TYPES)
         require("properties", properties, dict)
 
@@ -333,13 +461,78 @@ class Client(object):
             "event": "$pageview",
             "properties": properties,
             "timestamp": timestamp,
-            "context": context,
             "distinct_id": distinct_id,
             "uuid": uuid,
         }
 
         return self._enqueue(msg, disable_geoip)
 
+    def capture_exception(
+        self,
+        exception=None,
+        distinct_id=None,
+        properties=None,
+        context=None,
+        timestamp=None,
+        uuid=None,
+        groups=None,
+    ):
+        if context is not None:
+            warnings.warn(
+                "The 'context' parameter is deprecated and will be removed in a future version.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
+        # this function shouldn't ever throw an error, so it logs exceptions instead of raising them.
+        # this is important to ensure we don't unexpectedly re-raise exceptions in the user's code.
+        try:
+            properties = properties or {}
+
+            # if there's no distinct_id, we'll generate one and set personless mode
+            # via $process_person_profile = false
+            if distinct_id is None:
+                properties["$process_person_profile"] = False
+                distinct_id = uuid4()
+
+            require("distinct_id", distinct_id, ID_TYPES)
+            require("properties", properties, dict)
+
+            if exception is not None:
+                exc_info = exc_info_from_error(exception)
+            else:
+                exc_info = sys.exc_info()
+
+            if exc_info is None or exc_info == (None, None, None):
+                self.log.warning("No exception information available")
+                return
+
+            # Format stack trace for cymbal
+            all_exceptions_with_trace = exceptions_from_error_tuple(exc_info)
+
+            # Add in-app property to frames in the exceptions
+            event = handle_in_app(
+                {
+                    "exception": {
+                        "values": all_exceptions_with_trace,
+                    },
+                },
+                project_root=self.project_root,
+            )
+            all_exceptions_with_trace_and_in_app = event["exception"]["values"]
+
+            properties = {
+                "$exception_type": all_exceptions_with_trace_and_in_app[0].get("type"),
+                "$exception_message": all_exceptions_with_trace_and_in_app[0].get("value"),
+                "$exception_list": all_exceptions_with_trace_and_in_app,
+                "$exception_personURL": f"{remove_trailing_slash(self.raw_host)}/project/{self.api_key}/person/{distinct_id}",
+                **properties,
+            }
+
+            return self.capture(distinct_id, "$exception", properties, context, timestamp, uuid, groups)
+        except Exception as e:
+            self.log.exception(f"Failed to capture exception: {e}")
+
     def _enqueue(self, msg, disable_geoip):
         """Push a new `msg` onto the queue, return `(success, msg)`"""
 
@@ -348,10 +541,9 @@ class Client(object):
 
         timestamp = msg["timestamp"]
         if timestamp is None:
-            timestamp = datetime.utcnow().replace(tzinfo=tzutc())
+            timestamp = datetime.now(tz=tzutc())
 
         require("timestamp", timestamp, datetime)
-        require("context", msg["context"], dict)
 
         # add common
         timestamp = guess_timezone(timestamp)
@@ -374,6 +566,9 @@ class Client(object):
         if disable_geoip:
             msg["properties"]["$geoip_disable"] = True
 
+        if self.super_properties:
+            msg["properties"] = {**msg["properties"], **self.super_properties}
+
         msg["distinct_id"] = stringify_id(msg.get("distinct_id", None))
 
         msg = clean(msg)
@@ -432,6 +627,9 @@ class Client(object):
         self.flush()
         self.join()
 
+        if self.exception_capture:
+            self.exception_capture.close()
+
     def _load_feature_flags(self):
         try:
             response = get(
@@ -460,6 +658,19 @@ class Client(object):
                         "To use feature flags, please set a personal_api_key "
                         "More information: https://posthog.com/docs/api/overview",
                     )
+            elif e.status == 402:
+                self.log.warning("[FEATURE FLAGS] PostHog feature flags quota limited")
+                # Reset all feature flag data when quota limited
+                self.feature_flags = []
+                self.feature_flags_by_key = {}
+                self.group_type_mapping = {}
+                self.cohorts = {}
+
+                if self.debug:
+                    raise APIError(
+                        status=402,
+                        message="PostHog feature flags quota limited",
+                    )
             else:
                 self.log.error(f"[FEATURE FLAGS] Error loading feature flags: {e}")
         except Exception as e:
@@ -469,7 +680,7 @@ class Client(object):
             )
             self.log.warning(e)
 
-        self._last_feature_flag_poll = datetime.utcnow().replace(tzinfo=tzutc())
+        self._last_feature_flag_poll = datetime.now(tz=tzutc())
 
     def load_feature_flags(self):
         if not self.personal_api_key:
@@ -582,7 +793,6 @@ class Client(object):
             self.load_feature_flags()
         response = None
 
-        # If loading in previous line failed
         if self.feature_flags:
             for flag in self.feature_flags:
                 if flag["key"] == key:
@@ -601,6 +811,7 @@ class Client(object):
                     except Exception as e:
                         self.log.exception(f"[FEATURE FLAGS] Error while computing variant locally: {e}")
                         continue
+                    break
 
         flag_was_locally_evaluated = response is not None
         if not flag_was_locally_evaluated and not only_evaluate_locally:
@@ -662,23 +873,72 @@ class Client(object):
                 groups=groups,
                 person_properties=person_properties,
                 group_properties=group_properties,
-                send_feature_flag_events=send_feature_flag_events,
-                only_evaluate_locally=True,
+                send_feature_flag_events=False,
+                # Disable automatic sending of feature flag events because we're manually handling event dispatch.
+                # This prevents sending events with empty data when `get_feature_flag` cannot be evaluated locally.
+                only_evaluate_locally=True,  # Enable local evaluation of feature flags to avoid making multiple requests to `/decide`.
                 disable_geoip=disable_geoip,
             )
 
         response = None
+        payload = None
 
         if match_value is not None:
-            response = self._compute_payload_locally(key, match_value)
+            payload = self._compute_payload_locally(key, match_value)
 
-        if response is None and not only_evaluate_locally:
-            decide_payloads = self.get_feature_payloads(
-                distinct_id, groups, person_properties, group_properties, disable_geoip
+        flag_was_locally_evaluated = payload is not None
+        if not flag_was_locally_evaluated and not only_evaluate_locally:
+            try:
+                responses_and_payloads = self.get_feature_flags_and_payloads(
+                    distinct_id, groups, person_properties, group_properties, disable_geoip
+                )
+                response = responses_and_payloads["featureFlags"].get(key, None)
+                payload = responses_and_payloads["featureFlagPayloads"].get(str(key), None)
+            except Exception as e:
+                self.log.exception(f"[FEATURE FLAGS] Unable to get feature flags and payloads: {e}")
+
+        feature_flag_reported_key = f"{key}_{str(response)}"
+
+        if (
+            feature_flag_reported_key not in self.distinct_ids_feature_flags_reported[distinct_id]
+            and send_feature_flag_events  # noqa: W503
+        ):
+            self.capture(
+                distinct_id,
+                "$feature_flag_called",
+                {
+                    "$feature_flag": key,
+                    "$feature_flag_response": response,
+                    "$feature_flag_payload": payload,
+                    "locally_evaluated": flag_was_locally_evaluated,
+                    f"$feature/{key}": response,
+                },
+                groups=groups,
+                disable_geoip=disable_geoip,
             )
-            response = decide_payloads.get(str(key).lower(), None)
+            self.distinct_ids_feature_flags_reported[distinct_id].add(feature_flag_reported_key)
 
-        return response
+        return payload
+
+    def get_remote_config_payload(self, key: str):
+        if self.disabled:
+            return None
+
+        if self.personal_api_key is None:
+            self.log.warning(
+                "[FEATURE FLAGS] You have to specify a personal_api_key to fetch decrypted feature flag payloads."
+            )
+            return None
+
+        try:
+            return remote_config(
+                self.personal_api_key,
+                self.host,
+                key,
+                timeout=self.feature_flags_request_timeout_seconds,
+            )
+        except Exception as e:
+            self.log.exception(f"[FEATURE FLAGS] Unable to get decrypted feature flag payload: {e}")
 
     def _compute_payload_locally(self, key, match_value):
         payload = None
@@ -686,10 +946,14 @@ class Client(object):
         if self.feature_flags_by_key is None:
             return payload
 
-        flag_definition = self.feature_flags_by_key.get(key) or {}
-        flag_filters = flag_definition.get("filters") or {}
-        flag_payloads = flag_filters.get("payloads") or {}
-        payload = flag_payloads.get(str(match_value).lower(), None)
+        flag_definition = self.feature_flags_by_key.get(key)
+        if flag_definition:
+            flag_filters = flag_definition.get("filters") or {}
+            flag_payloads = flag_filters.get("payloads") or {}
+            # For boolean flags, convert True to "true"
+            # For multivariate flags, use the variant string as-is
+            lookup_value = "true" if isinstance(match_value, bool) and match_value else str(match_value)
+            payload = flag_payloads.get(lookup_value, None)
         return payload
 
     def get_all_flags(

posthog/exception_capture.py

@@ -1,53 +1,64 @@
 import logging
 import sys
 import threading
-from typing import TYPE_CHECKING
-
-from posthog.exception_utils import exceptions_from_error_tuple
+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"):
+    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.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)
+        self.capture_exception((args.exc_type, args.exc_value, args.exc_traceback))
 
-    def capture_exception(self, exc_type, exc_value, 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:
-            # if hasattr(sys, "ps1"):
-            #     # Disable the excepthook for interactive Python shells
-            #     return
-
-            # Format stack trace like sentry
-            all_exceptions_with_trace = exceptions_from_error_tuple((exc_type, exc_value, exc_traceback))
-
-            properties = {
-                "$exception_type": all_exceptions_with_trace[0].get("type"),
-                "$exception_message": all_exceptions_with_trace[0].get("value"),
-                "$exception_list": all_exceptions_with_trace,
-                # TODO: Can we somehow get distinct_id from context here? Stateless lib makes this much harder? 😅
-                # '$exception_personURL': f'{self.client.posthog_host}/project/{self.client.token}/person/{self.client.get_distinct_id()}'
-            }
-
-            # TODO: What distinct id should we attach these server-side exceptions to?
-            # Any heuristic seems prone to errors - how can we know if exception occurred in the context of a user that captured some other event?
-
-            self.client.capture("python-exceptions", "$exception", properties=properties)
+            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,88 @@
+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 {
+            "distinct_id": distinct_id,
+            "ip": headers.get("X-Forwarded-For"),
+            "user_agent": headers.get("User-Agent"),
+            "traceparent": traceparent,
+        }
+
+    def headers(self):
+        # type: () -> Dict[str, str]
+        return dict(self.request.headers)

posthog/exception_utils.py

@@ -369,6 +369,7 @@ def serialize_frame(
         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>",
@@ -417,7 +418,7 @@ def current_stacktrace(
 
     frames.reverse()
 
-    return {"frames": frames}
+    return {"frames": frames, "type": "raw"}
 
 
 def get_errno(exc_value):
@@ -503,7 +504,7 @@ def single_exception_from_error_tuple(
     ]
 
     if frames:
-        exception_value["stacktrace"] = {"frames": frames}
+        exception_value["stacktrace"] = {"frames": frames, "type": "raw"}
 
     return exception_value
 
@@ -792,7 +793,7 @@ def event_from_exception(
 
 
 def _module_in_list(name, items):
-    # type: (str, Optional[List[str]]) -> bool
+    # type: (str | None, Optional[List[str]]) -> bool
     if name is None:
         return False
 

posthog/feature_flags.py

@@ -7,6 +7,7 @@ from typing import Optional
 from dateutil import parser
 from dateutil.relativedelta import relativedelta
 
+from posthog import utils
 from posthog.utils import convert_to_datetime_aware, is_valid_regex
 
 __LONG_SCALE__ = float(0xFFFFFFFFFFFFFFF)
@@ -53,6 +54,9 @@ def match_feature_flag_properties(flag, distinct_id, properties, cohort_properti
     flag_conditions = (flag.get("filters") or {}).get("groups") or []
     is_inconclusive = False
     cohort_properties = cohort_properties or {}
+    # Some filters can be explicitly set to null, which require accessing variants like so
+    flag_variants = ((flag.get("filters") or {}).get("multivariate") or {}).get("variants") or []
+    valid_variant_keys = [variant["key"] for variant in flag_variants]
 
     # Stable sort conditions with variant overrides to the top. This ensures that if overrides are present, they are
     # evaluated first, and the variant override is applied to the first matching condition.
@@ -67,9 +71,7 @@ def match_feature_flag_properties(flag, distinct_id, properties, cohort_properti
             # the matching variant
             if is_condition_match(flag, distinct_id, condition, properties, cohort_properties):
                 variant_override = condition.get("variant")
-                # Some filters can be explicitly set to null, which require accessing variants like so
-                flag_variants = ((flag.get("filters") or {}).get("multivariate") or {}).get("variants") or []
-                if variant_override and variant_override in [variant["key"] for variant in flag_variants]:
+                if variant_override and variant_override in valid_variant_keys:
                     variant = variant_override
                 else:
                     variant = get_matching_variant(flag, distinct_id)
@@ -128,8 +130,8 @@ def match_property(property, property_values) -> bool:
 
         def compute_exact_match(value, override_value):
             if isinstance(value, list):
-                return str(override_value).lower() in [str(val).lower() for val in value]
-            return str(value).lower() == str(override_value).lower()
+                return str(override_value).casefold() in [str(val).casefold() for val in value]
+            return utils.str_iequals(value, override_value)
 
         if operator == "exact":
             return compute_exact_match(value, override_value)
@@ -140,10 +142,10 @@ def match_property(property, property_values) -> bool:
         return key in property_values
 
     if operator == "icontains":
-        return str(value).lower() in str(override_value).lower()
+        return utils.str_icontains(override_value, value)
 
     if operator == "not_icontains":
-        return str(value).lower() not in str(override_value).lower()
+        return not utils.str_icontains(override_value, value)
 
     if operator == "regex":
         return is_valid_regex(str(value)) and re.compile(str(value)).search(str(override_value)) is not None

posthog/request.py

@@ -11,7 +11,9 @@ from dateutil.tz import tzutc
 from posthog.utils import remove_trailing_slash
 from posthog.version import VERSION
 
+adapter = requests.adapters.HTTPAdapter(max_retries=2)
 _session = requests.sessions.Session()
+_session.mount("https://", adapter)
 
 US_INGESTION_ENDPOINT = "https://us.i.posthog.com"
 EU_INGESTION_ENDPOINT = "https://eu.i.posthog.com"
@@ -37,7 +39,7 @@ def post(
     """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)
@@ -66,7 +68,19 @@ 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("PostHog feature flags quota limited")
+            raise QuotaLimitError(res.status_code, "Feature flags quota limited")
+        return response
     try:
         payload = res.json()
         log.debug("received response: %s", payload)
@@ -81,6 +95,11 @@ def decide(api_key: str, host: Optional[str] = None, gzip: bool = False, timeout
     return _process_response(res, success_message="Feature flags decided successfully")
 
 
+def remote_config(personal_api_key: str, host: Optional[str] = None, key: str = "", timeout: int = 15) -> Any:
+    """Get remote config flag value from remote_config API endpoint"""
+    return get(personal_api_key, f"/api/projects/@current/feature_flags/{key}/remote_config/", host, timeout)
+
+
 def batch_post(
     api_key: str, host: Optional[str] = None, gzip: bool = False, timeout: int = 15, **kwargs
 ) -> requests.Response:
@@ -105,6 +124,10 @@ class APIError(Exception):
         return msg.format(self.message, self.status)
 
 
+class QuotaLimitError(APIError):
+    pass
+
+
 class DatetimeSerializer(json.JSONEncoder):
     def default(self, obj: Any):
         if isinstance(obj, (date, datetime)):

posthog/test/ai/anthropic/test_anthropic.py

@@ -0,0 +1,341 @@
+import os
+import time
+from unittest.mock import patch
+
+import pytest
+from anthropic.types import Message, Usage
+
+from posthog.ai.anthropic import Anthropic, AsyncAnthropic
+
+ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY")
+
+
+@pytest.fixture
+def mock_client():
+    with patch("posthog.client.Client") as mock_client:
+        mock_client.privacy_mode = False
+        yield mock_client
+
+
+@pytest.fixture
+def mock_anthropic_response():
+    return Message(
+        id="msg_123",
+        type="message",
+        role="assistant",
+        content=[{"type": "text", "text": "Test response"}],
+        model="claude-3-opus-20240229",
+        usage=Usage(
+            input_tokens=20,
+            output_tokens=10,
+        ),
+        stop_reason="end_turn",
+        stop_sequence=None,
+    )
+
+
+@pytest.fixture
+def mock_anthropic_stream():
+    class MockStreamEvent:
+        def __init__(self, content, usage=None):
+            self.content = content
+            self.usage = usage
+
+    def stream_generator():
+        yield MockStreamEvent("A")
+        yield MockStreamEvent("B")
+        yield MockStreamEvent(
+            "C",
+            usage=Usage(
+                input_tokens=20,
+                output_tokens=10,
+            ),
+        )
+
+    return stream_generator()
+
+
+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"

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,189 @@
+import time
+from unittest.mock import patch
+
+import pytest
+from openai.types.chat import ChatCompletion, ChatCompletionMessage
+from openai.types.chat.chat_completion import Choice
+from openai.types.completion_usage import CompletionUsage
+from openai.types.create_embedding_response import CreateEmbeddingResponse, Usage
+from openai.types.embedding import Embedding
+
+from posthog.ai.openai import OpenAI
+
+
+@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_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,
+        ),
+    )
+
+
+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"

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,68 @@
+from posthog.exception_integrations.django import DjangoRequestExtractor
+
+DEFAULT_USER_AGENT = (
+    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3"
+)
+
+
+def mock_request_factory(override_headers):
+    class Request:
+        META = {}
+        # TRICKY: Actual django request dict object has case insensitive matching, and strips http from the names
+        headers = {
+            "User-Agent": DEFAULT_USER_AGENT,
+            "Referrer": "http://example.com",
+            "X-Forwarded-For": "193.4.5.12",
+            **(override_headers or {}),
+        }
+
+    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,
+    }
+
+
+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,
+    }
+
+
+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",
+    }
+
+
+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",
+    }

posthog/test/test_client.py

@@ -5,8 +5,10 @@ from uuid import uuid4
 
 import mock
 import six
+from parameterized import parameterized
 
 from posthog.client import Client
+from posthog.request import APIError
 from posthog.test.test_utils import FAKE_TEST_API_KEY
 from posthog.version import VERSION
 
@@ -53,6 +55,11 @@ class TestClient(unittest.TestCase):
         self.assertEqual(msg["distinct_id"], "distinct_id")
         self.assertEqual(msg["properties"]["$lib"], "posthog-python")
         self.assertEqual(msg["properties"]["$lib_version"], VERSION)
+        # these will change between platforms so just asssert on presence here
+        assert msg["properties"]["$python_runtime"] == mock.ANY
+        assert msg["properties"]["$python_version"] == mock.ANY
+        assert msg["properties"]["$os"] == mock.ANY
+        assert msg["properties"]["$os_version"] == mock.ANY
 
     def test_basic_capture_with_uuid(self):
         client = self.client
@@ -84,6 +91,177 @@ class TestClient(unittest.TestCase):
         self.assertEqual(msg["properties"]["$lib"], "posthog-python")
         self.assertEqual(msg["properties"]["$lib_version"], VERSION)
 
+    def test_basic_super_properties(self):
+        client = Client(FAKE_TEST_API_KEY, super_properties={"source": "repo-name"})
+
+        _, msg = client.capture("distinct_id", "python test event")
+        client.flush()
+
+        self.assertEqual(msg["event"], "python test event")
+        self.assertEqual(msg["properties"]["source"], "repo-name")
+
+        _, msg = client.identify("distinct_id", {"trait": "value"})
+        client.flush()
+
+        self.assertEqual(msg["$set"]["trait"], "value")
+        self.assertEqual(msg["properties"]["source"], "repo-name")
+
+    def test_basic_capture_exception(self):
+        with mock.patch.object(Client, "capture", return_value=None) as patch_capture:
+            client = self.client
+            exception = Exception("test exception")
+            client.capture_exception(exception, distinct_id="distinct_id")
+
+            self.assertTrue(patch_capture.called)
+            capture_call = patch_capture.call_args[0]
+            self.assertEqual(capture_call[0], "distinct_id")
+            self.assertEqual(capture_call[1], "$exception")
+            self.assertEqual(
+                capture_call[2],
+                {
+                    "$exception_type": "Exception",
+                    "$exception_message": "test exception",
+                    "$exception_list": [
+                        {
+                            "mechanism": {"type": "generic", "handled": True},
+                            "module": None,
+                            "type": "Exception",
+                            "value": "test exception",
+                        }
+                    ],
+                    "$exception_personURL": "https://us.i.posthog.com/project/random_key/person/distinct_id",
+                },
+            )
+
+    def test_basic_capture_exception_with_distinct_id(self):
+        with mock.patch.object(Client, "capture", return_value=None) as patch_capture:
+            client = self.client
+            exception = Exception("test exception")
+            client.capture_exception(exception, "distinct_id")
+
+            self.assertTrue(patch_capture.called)
+            capture_call = patch_capture.call_args[0]
+            self.assertEqual(capture_call[0], "distinct_id")
+            self.assertEqual(capture_call[1], "$exception")
+            self.assertEqual(
+                capture_call[2],
+                {
+                    "$exception_type": "Exception",
+                    "$exception_message": "test exception",
+                    "$exception_list": [
+                        {
+                            "mechanism": {"type": "generic", "handled": True},
+                            "module": None,
+                            "type": "Exception",
+                            "value": "test exception",
+                        }
+                    ],
+                    "$exception_personURL": "https://us.i.posthog.com/project/random_key/person/distinct_id",
+                },
+            )
+
+    def test_basic_capture_exception_with_correct_host_generation(self):
+        with mock.patch.object(Client, "capture", return_value=None) as patch_capture:
+            client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, host="https://aloha.com")
+            exception = Exception("test exception")
+            client.capture_exception(exception, "distinct_id")
+
+            self.assertTrue(patch_capture.called)
+            capture_call = patch_capture.call_args[0]
+            self.assertEqual(capture_call[0], "distinct_id")
+            self.assertEqual(capture_call[1], "$exception")
+            self.assertEqual(
+                capture_call[2],
+                {
+                    "$exception_type": "Exception",
+                    "$exception_message": "test exception",
+                    "$exception_list": [
+                        {
+                            "mechanism": {"type": "generic", "handled": True},
+                            "module": None,
+                            "type": "Exception",
+                            "value": "test exception",
+                        }
+                    ],
+                    "$exception_personURL": "https://aloha.com/project/random_key/person/distinct_id",
+                },
+            )
+
+    def test_basic_capture_exception_with_correct_host_generation_for_server_hosts(self):
+        with mock.patch.object(Client, "capture", return_value=None) as patch_capture:
+            client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, host="https://app.posthog.com")
+            exception = Exception("test exception")
+            client.capture_exception(exception, "distinct_id")
+
+            self.assertTrue(patch_capture.called)
+            capture_call = patch_capture.call_args[0]
+            self.assertEqual(capture_call[0], "distinct_id")
+            self.assertEqual(capture_call[1], "$exception")
+            self.assertEqual(
+                capture_call[2],
+                {
+                    "$exception_type": "Exception",
+                    "$exception_message": "test exception",
+                    "$exception_list": [
+                        {
+                            "mechanism": {"type": "generic", "handled": True},
+                            "module": None,
+                            "type": "Exception",
+                            "value": "test exception",
+                        }
+                    ],
+                    "$exception_personURL": "https://app.posthog.com/project/random_key/person/distinct_id",
+                },
+            )
+
+    def test_basic_capture_exception_with_no_exception_given(self):
+        with mock.patch.object(Client, "capture", return_value=None) as patch_capture:
+            client = self.client
+            try:
+                raise Exception("test exception")
+            except Exception:
+                client.capture_exception(distinct_id="distinct_id")
+
+            self.assertTrue(patch_capture.called)
+            capture_call = patch_capture.call_args[0]
+            self.assertEqual(capture_call[0], "distinct_id")
+            self.assertEqual(capture_call[1], "$exception")
+            self.assertEqual(capture_call[2]["$exception_type"], "Exception")
+            self.assertEqual(capture_call[2]["$exception_message"], "test exception")
+            self.assertEqual(capture_call[2]["$exception_list"][0]["mechanism"]["type"], "generic")
+            self.assertEqual(capture_call[2]["$exception_list"][0]["mechanism"]["handled"], True)
+            self.assertEqual(capture_call[2]["$exception_list"][0]["module"], None)
+            self.assertEqual(capture_call[2]["$exception_list"][0]["type"], "Exception")
+            self.assertEqual(capture_call[2]["$exception_list"][0]["value"], "test exception")
+            self.assertEqual(
+                capture_call[2]["$exception_list"][0]["stacktrace"]["type"],
+                "raw",
+            )
+            self.assertEqual(
+                capture_call[2]["$exception_list"][0]["stacktrace"]["frames"][0]["filename"],
+                "posthog/test/test_client.py",
+            )
+            self.assertEqual(
+                capture_call[2]["$exception_list"][0]["stacktrace"]["frames"][0]["function"],
+                "test_basic_capture_exception_with_no_exception_given",
+            )
+            self.assertEqual(
+                capture_call[2]["$exception_list"][0]["stacktrace"]["frames"][0]["module"], "posthog.test.test_client"
+            )
+            self.assertEqual(capture_call[2]["$exception_list"][0]["stacktrace"]["frames"][0]["in_app"], True)
+
+    def test_basic_capture_exception_with_no_exception_happening(self):
+        with mock.patch.object(Client, "capture", return_value=None) as patch_capture:
+            with self.assertLogs("posthog", level="WARNING") as logs:
+                client = self.client
+                client.capture_exception()
+
+                self.assertFalse(patch_capture.called)
+                self.assertEqual(
+                    logs.output[0],
+                    "WARNING:posthog:No exception information available",
+                )
+
     @mock.patch("posthog.client.decide")
     def test_basic_capture_with_feature_flags(self, patch_decide):
         patch_decide.return_value = {"featureFlags": {"beta-feature": "random-variant"}}
@@ -114,7 +292,6 @@ class TestClient(unittest.TestCase):
             "id": 1,
             "name": "Beta Feature",
             "key": "beta-feature-local",
-            "is_simple_flag": False,
             "active": True,
             "rollout_percentage": 100,
             "filters": {
@@ -143,7 +320,6 @@ class TestClient(unittest.TestCase):
             "id": 1,
             "name": "Beta Feature",
             "key": "person-flag",
-            "is_simple_flag": True,
             "active": True,
             "filters": {
                 "groups": [
@@ -166,7 +342,6 @@ class TestClient(unittest.TestCase):
             "id": 1,
             "name": "Beta Feature",
             "key": "false-flag",
-            "is_simple_flag": True,
             "active": True,
             "filters": {
                 "groups": [
@@ -209,6 +384,25 @@ class TestClient(unittest.TestCase):
         assert "$feature/false-flag" not in msg["properties"]
         assert "$active_feature_flags" not in msg["properties"]
 
+    @mock.patch("posthog.client.get")
+    def test_load_feature_flags_quota_limited(self, patch_get):
+        mock_response = {
+            "type": "quota_limited",
+            "detail": "You have exceeded your feature flag request quota",
+            "code": "payment_required",
+        }
+        patch_get.side_effect = APIError(402, mock_response["detail"])
+
+        client = Client(FAKE_TEST_API_KEY, personal_api_key="test")
+        with self.assertLogs("posthog", level="WARNING") as logs:
+            client._load_feature_flags()
+
+            self.assertEqual(client.feature_flags, [])
+            self.assertEqual(client.feature_flags_by_key, {})
+            self.assertEqual(client.group_type_mapping, {})
+            self.assertEqual(client.cohorts, {})
+            self.assertIn("PostHog feature flags quota limited", logs.output[0])
+
     @mock.patch("posthog.client.decide")
     def test_dont_override_capture_with_local_flags(self, patch_decide):
         patch_decide.return_value = {"featureFlags": {"beta-feature": "random-variant"}}
@@ -218,7 +412,6 @@ class TestClient(unittest.TestCase):
             "id": 1,
             "name": "Beta Feature",
             "key": "beta-feature-local",
-            "is_simple_flag": False,
             "active": True,
             "rollout_percentage": 100,
             "filters": {
@@ -247,7 +440,6 @@ class TestClient(unittest.TestCase):
             "id": 1,
             "name": "Beta Feature",
             "key": "person-flag",
-            "is_simple_flag": True,
             "active": True,
             "filters": {
                 "groups": [
@@ -403,16 +595,14 @@ class TestClient(unittest.TestCase):
             "distinct_id",
             "python test event",
             {"property": "value"},
-            {"ip": "192.168.0.1"},
-            datetime(2014, 9, 3),
-            "new-uuid",
+            timestamp=datetime(2014, 9, 3),
+            uuid="new-uuid",
         )
 
         self.assertTrue(success)
 
         self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00")
         self.assertEqual(msg["properties"]["property"], "value")
-        self.assertEqual(msg["context"]["ip"], "192.168.0.1")
         self.assertEqual(msg["event"], "python test event")
         self.assertEqual(msg["properties"]["$lib"], "posthog-python")
         self.assertEqual(msg["properties"]["$lib_version"], VERSION)
@@ -445,13 +635,12 @@ class TestClient(unittest.TestCase):
     def test_advanced_identify(self):
         client = self.client
         success, msg = client.identify(
-            "distinct_id", {"trait": "value"}, {"ip": "192.168.0.1"}, datetime(2014, 9, 3), "new-uuid"
+            "distinct_id", {"trait": "value"}, timestamp=datetime(2014, 9, 3), uuid="new-uuid"
         )
 
         self.assertTrue(success)
 
         self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00")
-        self.assertEqual(msg["context"]["ip"], "192.168.0.1")
         self.assertEqual(msg["$set"]["trait"], "value")
         self.assertEqual(msg["properties"]["$lib"], "posthog-python")
         self.assertEqual(msg["properties"]["$lib_version"], VERSION)
@@ -473,14 +662,11 @@ class TestClient(unittest.TestCase):
 
     def test_advanced_set(self):
         client = self.client
-        success, msg = client.set(
-            "distinct_id", {"trait": "value"}, {"ip": "192.168.0.1"}, datetime(2014, 9, 3), "new-uuid"
-        )
+        success, msg = client.set("distinct_id", {"trait": "value"}, timestamp=datetime(2014, 9, 3), uuid="new-uuid")
 
         self.assertTrue(success)
 
         self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00")
-        self.assertEqual(msg["context"]["ip"], "192.168.0.1")
         self.assertEqual(msg["$set"]["trait"], "value")
         self.assertEqual(msg["properties"]["$lib"], "posthog-python")
         self.assertEqual(msg["properties"]["$lib_version"], VERSION)
@@ -503,13 +689,12 @@ class TestClient(unittest.TestCase):
     def test_advanced_set_once(self):
         client = self.client
         success, msg = client.set_once(
-            "distinct_id", {"trait": "value"}, {"ip": "192.168.0.1"}, datetime(2014, 9, 3), "new-uuid"
+            "distinct_id", {"trait": "value"}, timestamp=datetime(2014, 9, 3), uuid="new-uuid"
         )
 
         self.assertTrue(success)
 
         self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00")
-        self.assertEqual(msg["context"]["ip"], "192.168.0.1")
         self.assertEqual(msg["$set_once"]["trait"], "value")
         self.assertEqual(msg["properties"]["$lib"], "posthog-python")
         self.assertEqual(msg["properties"]["$lib_version"], VERSION)
@@ -537,9 +722,28 @@ class TestClient(unittest.TestCase):
         self.assertTrue(isinstance(msg["timestamp"], str))
         self.assertIsNone(msg.get("uuid"))
 
+    def test_basic_group_identify_with_distinct_id(self):
+        success, msg = self.client.group_identify("organization", "id:5", distinct_id="distinct_id")
+        self.assertTrue(success)
+        self.assertEqual(msg["event"], "$groupidentify")
+        self.assertEqual(msg["distinct_id"], "distinct_id")
+        self.assertEqual(
+            msg["properties"],
+            {
+                "$group_type": "organization",
+                "$group_key": "id:5",
+                "$group_set": {},
+                "$lib": "posthog-python",
+                "$lib_version": VERSION,
+                "$geoip_disable": True,
+            },
+        )
+        self.assertTrue(isinstance(msg["timestamp"], str))
+        self.assertIsNone(msg.get("uuid"))
+
     def test_advanced_group_identify(self):
         success, msg = self.client.group_identify(
-            "organization", "id:5", {"trait": "value"}, {"ip": "192.168.0.1"}, datetime(2014, 9, 3), "new-uuid"
+            "organization", "id:5", {"trait": "value"}, timestamp=datetime(2014, 9, 3), uuid="new-uuid"
         )
 
         self.assertTrue(success)
@@ -557,7 +761,33 @@ class TestClient(unittest.TestCase):
             },
         )
         self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00")
-        self.assertEqual(msg["context"]["ip"], "192.168.0.1")
+
+    def test_advanced_group_identify_with_distinct_id(self):
+        success, msg = self.client.group_identify(
+            "organization",
+            "id:5",
+            {"trait": "value"},
+            timestamp=datetime(2014, 9, 3),
+            uuid="new-uuid",
+            distinct_id="distinct_id",
+        )
+
+        self.assertTrue(success)
+        self.assertEqual(msg["event"], "$groupidentify")
+        self.assertEqual(msg["distinct_id"], "distinct_id")
+
+        self.assertEqual(
+            msg["properties"],
+            {
+                "$group_type": "organization",
+                "$group_key": "id:5",
+                "$group_set": {"trait": "value"},
+                "$lib": "posthog-python",
+                "$lib_version": VERSION,
+                "$geoip_disable": True,
+            },
+        )
+        self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00")
 
     def test_basic_alias(self):
         client = self.client
@@ -593,15 +823,13 @@ class TestClient(unittest.TestCase):
             "distinct_id",
             "https://posthog.com/contact",
             {"property": "value"},
-            {"ip": "192.168.0.1"},
-            datetime(2014, 9, 3),
-            "new-uuid",
+            timestamp=datetime(2014, 9, 3),
+            uuid="new-uuid",
         )
 
         self.assertTrue(success)
 
         self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00")
-        self.assertEqual(msg["context"]["ip"], "192.168.0.1")
         self.assertEqual(msg["properties"]["$current_url"], "https://posthog.com/contact")
         self.assertEqual(msg["properties"]["property"], "value")
         self.assertEqual(msg["properties"]["$lib"], "posthog-python")
@@ -825,7 +1053,7 @@ class TestClient(unittest.TestCase):
 
         patch_get.return_value.raiseError.side_effect = raise_effect
         client = Client(FAKE_TEST_API_KEY, personal_api_key="test")
-        client.feature_flags = [{"key": "example", "is_simple_flag": False}]
+        client.feature_flags = [{"key": "example"}]
 
         self.assertFalse(client.feature_enabled("example", "distinct_id"))
 
@@ -895,3 +1123,92 @@ class TestClient(unittest.TestCase):
             group_properties={},
             disable_geoip=False,
         )
+
+    @parameterized.expand(
+        [
+            # name, sys_platform, version_info, expected_runtime, expected_version, expected_os, expected_os_version, platform_method, platform_return, distro_info
+            (
+                "macOS",
+                "darwin",
+                (3, 8, 10),
+                "MockPython",
+                "3.8.10",
+                "Mac OS X",
+                "10.15.7",
+                "mac_ver",
+                ("10.15.7", "", ""),
+                None,
+            ),
+            (
+                "Windows",
+                "win32",
+                (3, 8, 10),
+                "MockPython",
+                "3.8.10",
+                "Windows",
+                "10",
+                "win32_ver",
+                ("10", "", "", ""),
+                None,
+            ),
+            (
+                "Linux",
+                "linux",
+                (3, 8, 10),
+                "MockPython",
+                "3.8.10",
+                "Linux",
+                "20.04",
+                None,
+                None,
+                {"version": "20.04"},
+            ),
+        ]
+    )
+    def test_mock_system_context(
+        self,
+        _name,
+        sys_platform,
+        version_info,
+        expected_runtime,
+        expected_version,
+        expected_os,
+        expected_os_version,
+        platform_method,
+        platform_return,
+        distro_info,
+    ):
+        """Test that we can mock platform and sys for testing system_context"""
+        with mock.patch("posthog.client.platform") as mock_platform:
+            with mock.patch("posthog.client.sys") as mock_sys:
+                # Set up common mocks
+                mock_platform.python_implementation.return_value = expected_runtime
+                mock_sys.version_info = version_info
+                mock_sys.platform = sys_platform
+
+                # Set up platform-specific mocks
+                if platform_method:
+                    getattr(mock_platform, platform_method).return_value = platform_return
+
+                # Special handling for Linux which uses distro module
+                if sys_platform == "linux":
+                    # Directly patch the get_os_info function to return our expected values
+                    with mock.patch("posthog.client.get_os_info", return_value=(expected_os, expected_os_version)):
+                        from posthog.client import system_context
+
+                        context = system_context()
+                else:
+                    # Get system context for non-Linux platforms
+                    from posthog.client import system_context
+
+                    context = system_context()
+
+                # Verify results
+                expected_context = {
+                    "$python_runtime": expected_runtime,
+                    "$python_version": expected_version,
+                    "$os": expected_os,
+                    "$os_version": expected_os_version,
+                }
+
+                assert context == expected_context

posthog/test/test_exception_capture.py

@@ -29,6 +29,35 @@ def test_excepthook(tmpdir):
     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": [{"filename": "app.py", "abs_path"'
+        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_flags.py

@@ -38,7 +38,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "person-flag",
-                "is_simple_flag": True,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -69,6 +68,59 @@ class TestLocalEvaluation(unittest.TestCase):
         self.assertTrue(feature_flag_match)
         self.assertFalse(not_feature_flag_match)
 
+    def test_case_insensitive_matching(self):
+        self.client.feature_flags = [
+            {
+                "id": 1,
+                "name": "Beta Feature",
+                "key": "person-flag",
+                "is_simple_flag": True,
+                "active": True,
+                "filters": {
+                    "groups": [
+                        {
+                            "properties": [
+                                {
+                                    "key": "location",
+                                    "operator": "exact",
+                                    "value": ["Straße"],
+                                    "type": "person",
+                                }
+                            ],
+                            "rollout_percentage": 100,
+                        },
+                        {
+                            "properties": [
+                                {
+                                    "key": "star",
+                                    "operator": "exact",
+                                    "value": ["ſun"],
+                                    "type": "person",
+                                }
+                            ],
+                            "rollout_percentage": 100,
+                        },
+                    ],
+                },
+            }
+        ]
+
+        self.assertTrue(
+            self.client.get_feature_flag("person-flag", "some-distinct-id", person_properties={"location": "straße"})
+        )
+
+        self.assertTrue(
+            self.client.get_feature_flag("person-flag", "some-distinct-id", person_properties={"location": "strasse"})
+        )
+
+        self.assertTrue(
+            self.client.get_feature_flag("person-flag", "some-distinct-id", person_properties={"star": "ſun"})
+        )
+
+        self.assertTrue(
+            self.client.get_feature_flag("person-flag", "some-distinct-id", person_properties={"star": "sun"})
+        )
+
     @mock.patch("posthog.client.decide")
     @mock.patch("posthog.client.get")
     def test_flag_group_properties(self, patch_get, patch_decide):
@@ -77,7 +129,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "group-flag",
-                "is_simple_flag": True,
                 "active": True,
                 "filters": {
                     "aggregation_group_type_index": 0,
@@ -170,7 +221,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "complex-flag",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -286,7 +336,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": True,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -301,7 +350,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 2,
                 "name": "Beta Feature",
                 "key": "beta-feature2",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -343,7 +391,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": True,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -358,7 +405,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 2,
                 "name": "Beta Feature",
                 "key": "beta-feature2",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -410,7 +456,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": True,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -457,7 +502,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "rollout_percentage": 100,
                 "filters": {
@@ -472,7 +516,7 @@ class TestLocalEvaluation(unittest.TestCase):
             }
         ]
         # decide called always because experience_continuity is set
-        self.assertTrue(client.get_feature_flag("beta-feature", "distinct_id"), "decide-fallback-value")
+        self.assertEqual(client.get_feature_flag("beta-feature", "distinct_id"), "decide-fallback-value")
         self.assertEqual(patch_decide.call_count, 1)
 
     @mock.patch.object(Client, "capture")
@@ -487,7 +531,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "rollout_percentage": 100,
                 "filters": {
@@ -503,7 +546,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 2,
                 "name": "Beta Feature",
                 "key": "disabled-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -518,7 +560,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 3,
                 "name": "Beta Feature",
                 "key": "beta-feature2",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -551,7 +592,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "rollout_percentage": 100,
                 "filters": {
@@ -570,7 +610,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 2,
                 "name": "Beta Feature",
                 "key": "disabled-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -588,7 +627,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 3,
                 "name": "Beta Feature",
                 "key": "beta-feature2",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -654,7 +692,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "rollout_percentage": 100,
                 "filters": {
@@ -670,7 +707,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 2,
                 "name": "Beta Feature",
                 "key": "disabled-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -696,7 +732,6 @@ class TestLocalEvaluation(unittest.TestCase):
             "id": 1,
             "name": "Beta Feature",
             "key": "beta-feature",
-            "is_simple_flag": False,
             "active": True,
             "rollout_percentage": 100,
             "filters": {
@@ -715,7 +750,6 @@ class TestLocalEvaluation(unittest.TestCase):
             "id": 2,
             "name": "Beta Feature",
             "key": "disabled-feature",
-            "is_simple_flag": False,
             "active": True,
             "filters": {
                 "groups": [
@@ -751,7 +785,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "rollout_percentage": 100,
                 "filters": {
@@ -767,7 +800,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 2,
                 "name": "Beta Feature",
                 "key": "disabled-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -782,7 +814,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 3,
                 "name": "Beta Feature",
                 "key": "beta-feature2",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -814,7 +845,6 @@ class TestLocalEvaluation(unittest.TestCase):
             "id": 1,
             "name": "Beta Feature",
             "key": "beta-feature",
-            "is_simple_flag": False,
             "active": True,
             "rollout_percentage": 100,
             "filters": {
@@ -833,7 +863,6 @@ class TestLocalEvaluation(unittest.TestCase):
             "id": 2,
             "name": "Beta Feature",
             "key": "disabled-feature",
-            "is_simple_flag": False,
             "active": True,
             "filters": {
                 "groups": [
@@ -851,7 +880,6 @@ class TestLocalEvaluation(unittest.TestCase):
             "id": 3,
             "name": "Beta Feature",
             "key": "beta-feature2",
-            "is_simple_flag": False,
             "active": True,
             "filters": {
                 "groups": [
@@ -888,7 +916,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "rollout_percentage": 100,
                 "filters": {
@@ -904,7 +931,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 2,
                 "name": "Beta Feature",
                 "key": "disabled-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -927,7 +953,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": False,
                 "active": False,
                 "rollout_percentage": 100,
                 "filters": {
@@ -943,7 +968,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 2,
                 "name": "Beta Feature",
                 "key": "disabled-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -969,7 +993,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 id: 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": True,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -1007,13 +1030,15 @@ class TestLocalEvaluation(unittest.TestCase):
             "beta-feature",
             "some-distinct-id",
             person_properties={
-                "latestBuildVersion": "24.32..1",
+                "latestBuildVersion": "24.32.1",
                 "latestBuildVersionMajor": "24",
                 "latestBuildVersionMinor": "32",
                 "latestBuildVersionPatch": "1",
             },
         )
 
+        self.assertEqual(feature_flag_match, True)
+
     @mock.patch("posthog.client.decide")
     @mock.patch("posthog.client.get")
     def test_feature_flags_local_evaluation_for_cohorts(self, patch_get, patch_decide):
@@ -1023,7 +1048,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 2,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -1094,7 +1118,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 2,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -1207,7 +1230,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": True,
                 "active": True,
                 "rollout_percentage": 100,
                 "filters": {
@@ -1232,7 +1254,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": True,
                 "active": True,
                 "rollout_percentage": 0,
                 "filters": {
@@ -1257,7 +1278,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": True,
                 "active": True,
                 "rollout_percentage": None,
                 "filters": {
@@ -1281,7 +1301,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": True,
                 "active": True,
                 "rollout_percentage": 100,
                 "filters": {
@@ -1305,7 +1324,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "rollout_percentage": 100,
                 "filters": {
@@ -1330,7 +1348,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": True,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -1352,7 +1369,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "rollout_percentage": 100,
                 "filters": {
@@ -1418,7 +1434,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "rollout_percentage": 100,
                 "filters": {
@@ -1459,7 +1474,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "rollout_percentage": 100,
                 "filters": {
@@ -1511,7 +1525,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "rollout_percentage": 100,
                 "filters": {
@@ -1552,7 +1565,6 @@ class TestLocalEvaluation(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "beta-feature",
-                "is_simple_flag": False,
                 "active": True,
                 "rollout_percentage": 100,
                 "filters": {
@@ -1595,7 +1607,6 @@ class TestLocalEvaluation(unittest.TestCase):
             "id": 1,
             "name": "Beta Feature",
             "key": "person-flag",
-            "is_simple_flag": True,
             "active": True,
             "filters": {
                 "groups": [
@@ -1632,9 +1643,10 @@ class TestLocalEvaluation(unittest.TestCase):
         )
         self.assertEqual(patch_decide.call_count, 0)
 
+    @mock.patch.object(Client, "capture")
     @mock.patch("posthog.client.decide")
-    def test_boolean_feature_flag_payload_decide(self, patch_decide):
-        patch_decide.return_value = {"featureFlagPayloads": {"person-flag": 300}}
+    def test_boolean_feature_flag_payload_decide(self, patch_decide, patch_capture):
+        patch_decide.return_value = {"featureFlags": {"person-flag": True}, "featureFlagPayloads": {"person-flag": 300}}
         self.assertEqual(
             self.client.get_feature_flag_payload(
                 "person-flag", "some-distinct-id", person_properties={"region": "USA"}
@@ -1649,6 +1661,8 @@ class TestLocalEvaluation(unittest.TestCase):
             300,
         )
         self.assertEqual(patch_decide.call_count, 2)
+        self.assertEqual(patch_capture.call_count, 1)
+        patch_capture.reset_mock()
 
     @mock.patch("posthog.client.decide")
     def test_multivariate_feature_flag_payloads(self, patch_decide):
@@ -1656,7 +1670,6 @@ class TestLocalEvaluation(unittest.TestCase):
             "id": 1,
             "name": "Beta Feature",
             "key": "beta-feature",
-            "is_simple_flag": False,
             "active": True,
             "rollout_percentage": 100,
             "filters": {
@@ -2234,7 +2247,6 @@ class TestCaptureCalls(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "complex-flag",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -2334,6 +2346,136 @@ class TestCaptureCalls(unittest.TestCase):
             disable_geoip=None,
         )
 
+    @mock.patch("posthog.client.decide")
+    def test_capture_is_called_but_does_not_add_all_flags(self, patch_decide):
+        patch_decide.return_value = {"featureFlags": {"decide-flag": "decide-value"}}
+        client = Client(FAKE_TEST_API_KEY, personal_api_key=FAKE_TEST_API_KEY)
+        client.feature_flags = [
+            {
+                "id": 1,
+                "name": "Beta Feature",
+                "key": "complex-flag",
+                "active": True,
+                "filters": {
+                    "groups": [
+                        {
+                            "properties": [{"key": "region", "value": "USA"}],
+                            "rollout_percentage": 100,
+                        },
+                    ],
+                },
+            },
+            {
+                "id": 2,
+                "name": "Gamma Feature",
+                "key": "simple-flag",
+                "active": True,
+                "filters": {
+                    "groups": [
+                        {
+                            "properties": [],
+                            "rollout_percentage": 100,
+                        },
+                    ],
+                },
+            },
+        ]
+
+        self.assertTrue(
+            client.get_feature_flag("complex-flag", "some-distinct-id", person_properties={"region": "USA"})
+        )
+
+        # Grab the capture message that was just added to the queue
+        msg = client.queue.get(block=False)
+        assert msg["event"] == "$feature_flag_called"
+        assert msg["properties"]["$feature_flag"] == "complex-flag"
+        assert msg["properties"]["$feature_flag_response"] is True
+        assert msg["properties"]["locally_evaluated"] is True
+        assert msg["properties"]["$feature/complex-flag"] is True
+        assert "$feature/simple-flag" not in msg["properties"]
+        assert "$active_feature_flags" not in msg["properties"]
+
+    @mock.patch.object(Client, "capture")
+    @mock.patch("posthog.client.decide")
+    def test_capture_is_called_in_get_feature_flag_payload(self, patch_decide, patch_capture):
+        patch_decide.return_value = {
+            "featureFlags": {"person-flag": True},
+            "featureFlagPayloads": {"person-flag": 300},
+        }
+        client = Client(api_key=FAKE_TEST_API_KEY, personal_api_key=FAKE_TEST_API_KEY)
+
+        client.feature_flags = [
+            {
+                "id": 1,
+                "name": "Beta Feature",
+                "key": "person-flag",
+                "active": True,
+                "filters": {
+                    "groups": [
+                        {
+                            "properties": [{"key": "region", "value": "USA"}],
+                            "rollout_percentage": 100,
+                        }
+                    ],
+                },
+            }
+        ]
+
+        # Call get_feature_flag_payload with match_value=None to trigger get_feature_flag
+        client.get_feature_flag_payload(
+            key="person-flag", distinct_id="some-distinct-id", person_properties={"region": "USA", "name": "Aloha"}
+        )
+
+        # Assert that capture was called once, with the correct parameters
+        self.assertEqual(patch_capture.call_count, 1)
+        patch_capture.assert_called_with(
+            "some-distinct-id",
+            "$feature_flag_called",
+            {
+                "$feature_flag": "person-flag",
+                "$feature_flag_response": True,
+                "$feature_flag_payload": 300,
+                "locally_evaluated": False,
+                "$feature/person-flag": True,
+            },
+            groups={},
+            disable_geoip=None,
+        )
+
+        # Reset mocks for further tests
+        patch_capture.reset_mock()
+        patch_decide.reset_mock()
+
+        # Call get_feature_flag_payload again for the same user; capture should not be called again because we've already reported an event for this distinct_id + flag
+        client.get_feature_flag_payload(
+            key="person-flag", distinct_id="some-distinct-id", person_properties={"region": "USA", "name": "Aloha"}
+        )
+
+        self.assertEqual(patch_capture.call_count, 0)
+        patch_capture.reset_mock()
+
+        # Call get_feature_flag_payload for a different user; capture should be called
+        client.get_feature_flag_payload(
+            key="person-flag", distinct_id="some-distinct-id2", person_properties={"region": "USA", "name": "Aloha"}
+        )
+
+        self.assertEqual(patch_capture.call_count, 1)
+        patch_capture.assert_called_with(
+            "some-distinct-id2",
+            "$feature_flag_called",
+            {
+                "$feature_flag": "person-flag",
+                "$feature_flag_response": True,
+                "$feature_flag_payload": 300,
+                "locally_evaluated": False,
+                "$feature/person-flag": True,
+            },
+            groups={},
+            disable_geoip=None,
+        )
+
+        patch_capture.reset_mock()
+
     @mock.patch.object(Client, "capture")
     @mock.patch("posthog.client.decide")
     def test_disable_geoip_get_flag_capture_call(self, patch_decide, patch_capture):
@@ -2344,7 +2486,6 @@ class TestCaptureCalls(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "complex-flag",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -2387,7 +2528,6 @@ class TestCaptureCalls(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "complex-flag",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [
@@ -2446,7 +2586,6 @@ class TestConsistency(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "simple-flag",
-                "is_simple_flag": True,
                 "active": True,
                 "filters": {
                     "groups": [{"properties": [], "rollout_percentage": 45}],
@@ -3474,7 +3613,6 @@ class TestConsistency(unittest.TestCase):
                 "id": 1,
                 "name": "Beta Feature",
                 "key": "multivariate-flag",
-                "is_simple_flag": False,
                 "active": True,
                 "filters": {
                     "groups": [{"properties": [], "rollout_percentage": 55}],
@@ -4502,3 +4640,84 @@ class TestConsistency(unittest.TestCase):
                 self.assertEqual(feature_flag_match, results[i])
             else:
                 self.assertFalse(feature_flag_match)
+
+    @mock.patch("posthog.client.decide")
+    def test_feature_flag_case_sensitive(self, mock_decide):
+        mock_decide.return_value = {"featureFlags": {}}  # Ensure decide returns empty flags
+
+        client = Client(api_key=FAKE_TEST_API_KEY, personal_api_key=FAKE_TEST_API_KEY)
+        client.feature_flags = [
+            {
+                "id": 1,
+                "key": "Beta-Feature",
+                "active": True,
+                "filters": {
+                    "groups": [{"properties": [], "rollout_percentage": 100}],
+                },
+            }
+        ]
+
+        # Test that flag evaluation is case-sensitive
+        self.assertTrue(client.feature_enabled("Beta-Feature", "user1"))
+        self.assertFalse(client.feature_enabled("beta-feature", "user1"))
+        self.assertFalse(client.feature_enabled("BETA-FEATURE", "user1"))
+
+    @mock.patch("posthog.client.decide")
+    def test_feature_flag_payload_case_sensitive(self, mock_decide):
+        mock_decide.return_value = {
+            "featureFlags": {"Beta-Feature": True},
+            "featureFlagPayloads": {"Beta-Feature": {"some": "value"}},
+        }
+
+        client = Client(api_key=FAKE_TEST_API_KEY, personal_api_key=FAKE_TEST_API_KEY)
+        client.feature_flags = [
+            {
+                "id": 1,
+                "key": "Beta-Feature",
+                "active": True,
+                "filters": {
+                    "groups": [{"properties": [], "rollout_percentage": 100}],
+                    "payloads": {
+                        "true": {"some": "value"},
+                    },
+                },
+            }
+        ]
+
+        # Test that payload retrieval is case-sensitive
+        self.assertEqual(client.get_feature_flag_payload("Beta-Feature", "user1"), {"some": "value"})
+        self.assertIsNone(client.get_feature_flag_payload("beta-feature", "user1"))
+        self.assertIsNone(client.get_feature_flag_payload("BETA-FEATURE", "user1"))
+
+    @mock.patch("posthog.client.decide")
+    def test_feature_flag_case_sensitive_consistency(self, mock_decide):
+        mock_decide.return_value = {
+            "featureFlags": {"Beta-Feature": True},
+            "featureFlagPayloads": {"Beta-Feature": {"some": "value"}},
+        }
+
+        client = Client(api_key=FAKE_TEST_API_KEY, personal_api_key=FAKE_TEST_API_KEY)
+        client.feature_flags = [
+            {
+                "id": 1,
+                "key": "Beta-Feature",
+                "active": True,
+                "filters": {
+                    "groups": [{"properties": [], "rollout_percentage": 100}],
+                    "payloads": {
+                        "true": {"some": "value"},
+                    },
+                },
+            }
+        ]
+
+        # Test that flag evaluation and payload retrieval are consistently case-sensitive
+        # Only exact match should work
+        self.assertTrue(client.feature_enabled("Beta-Feature", "user1"))
+        self.assertEqual(client.get_feature_flag_payload("Beta-Feature", "user1"), {"some": "value"})
+
+        # Different cases should not match
+        test_cases = ["beta-feature", "BETA-FEATURE", "bEtA-FeAtUrE"]
+        for case in test_cases:
+            self.assertFalse(client.feature_enabled(case, "user1"))
+            self.assertIsNone(client.get_feature_flag_payload(case, "user1"))

posthog/test/test_request.py

@@ -2,10 +2,11 @@ import json
 import unittest
 from datetime import date, datetime
 
+import mock
 import pytest
 import requests
 
-from posthog.request import DatetimeSerializer, batch_post, determine_server_host
+from posthog.request import DatetimeSerializer, QuotaLimitError, batch_post, decide, determine_server_host
 from posthog.test.test_utils import TEST_API_KEY
 
 
@@ -44,6 +45,36 @@ class TestRequests(unittest.TestCase):
                 "key", batch=[{"distinct_id": "distinct_id", "event": "python event", "type": "track"}], timeout=0.0001
             )
 
+    def test_quota_limited_response(self):
+        mock_response = requests.Response()
+        mock_response.status_code = 200
+        mock_response._content = json.dumps(
+            {
+                "quotaLimited": ["feature_flags"],
+                "featureFlags": {},
+                "featureFlagPayloads": {},
+                "errorsWhileComputingFlags": False,
+            }
+        ).encode("utf-8")
+
+        with mock.patch("posthog.request._session.post", return_value=mock_response):
+            with self.assertRaises(QuotaLimitError) as cm:
+                decide("fake_key", "fake_host")
+
+            self.assertEqual(cm.exception.status, 200)
+            self.assertEqual(cm.exception.message, "Feature flags quota limited")
+
+    def test_normal_decide_response(self):
+        mock_response = requests.Response()
+        mock_response.status_code = 200
+        mock_response._content = json.dumps(
+            {"featureFlags": {"flag1": True}, "featureFlagPayloads": {}, "errorsWhileComputingFlags": False}
+        ).encode("utf-8")
+
+        with mock.patch("posthog.request._session.post", return_value=mock_response):
+            response = decide("fake_key", "fake_host")
+            self.assertEqual(response["featureFlags"], {"flag1": True})
+
 
 @pytest.mark.parametrize(
     "host, expected",

posthog/test/test_utils.py

@@ -1,10 +1,13 @@
 import unittest
 from datetime import date, datetime, timedelta
 from decimal import Decimal
+from typing import Optional
 from uuid import UUID
 
 import six
 from dateutil.tz import tzutc
+from pydantic import BaseModel
+from pydantic.v1 import BaseModel as BaseModelV1
 
 from posthog import utils
 
@@ -81,6 +84,32 @@ class TestUtils(unittest.TestCase):
         self.assertEqual("http://posthog.io", utils.remove_trailing_slash("http://posthog.io/"))
         self.assertEqual("http://posthog.io", utils.remove_trailing_slash("http://posthog.io"))
 
+    def test_clean_pydantic(self):
+        class ModelV2(BaseModel):
+            foo: str
+            bar: int
+            baz: Optional[str] = None
+
+        class ModelV1(BaseModelV1):
+            foo: int
+            bar: str
+
+        class NestedModel(BaseModel):
+            foo: ModelV2
+
+        self.assertEqual(utils.clean(ModelV2(foo="1", bar=2)), {"foo": "1", "bar": 2, "baz": None})
+        self.assertEqual(utils.clean(ModelV1(foo=1, bar="2")), {"foo": 1, "bar": "2"})
+        self.assertEqual(
+            utils.clean(NestedModel(foo=ModelV2(foo="1", bar=2, baz="3"))), {"foo": {"foo": "1", "bar": 2, "baz": "3"}}
+        )
+
+        class Dummy:
+            def model_dump(self, required_param):
+                pass
+
+        # Skips a class with a defined non-Pydantic `model_dump` method.
+        self.assertEqual(utils.clean({"test": Dummy()}), {})
+
 
 class TestSizeLimitedDict(unittest.TestCase):
     def test_size_limited_dict(self):

posthog/utils.py

@@ -51,14 +51,24 @@ 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)
+    return _coerce_unicode(item)
 
 
 def _clean_list(list_):
@@ -115,3 +125,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.6.0"
+VERSION = "3.18.0"
 
 if __name__ == "__main__":
     print(VERSION, end="")  # noqa: T201

setup.cfg

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

setup.py

@@ -14,18 +14,54 @@ long_description = """
 PostHog is developer-friendly, self-hosted product analytics. posthog-python is the python package.
 """
 
-install_requires = ["requests>=2.7,<3.0", "six>=1.5", "monotonic>=1.5", "backoff>=1.10.0", "python-dateutil>2.1"]
+install_requires = [
+    "requests>=2.7,<3.0",
+    "six>=1.5",
+    "monotonic>=1.5",
+    "backoff>=1.10.0",
+    "python-dateutil>2.1",
+    "distro>=1.5.0",  # Required for Linux OS detection in Python 3.9+
+]
 
 extras_require = {
     "dev": [
         "black",
+        "django-stubs",
         "isort",
         "flake8",
         "flake8-print",
+        "lxml",
+        "mypy",
+        "mypy-baseline",
+        "types-mock",
+        "types-python-dateutil",
+        "types-requests",
+        "types-setuptools",
+        "types-six",
         "pre-commit",
+        "pydantic",
+    ],
+    "test": [
+        "mock>=2.0.0",
+        "freezegun==1.5.1",
+        "pylint",
+        "flake8",
+        "coverage",
+        "pytest",
+        "pytest-timeout",
+        "pytest-asyncio",
+        "django",
+        "openai",
+        "anthropic",
+        "langgraph",
+        "langchain-community>=0.2.0",
+        "langchain-openai>=0.2.0",
+        "langchain-anthropic>=0.2.0",
+        "pydantic",
+        "parameterized>=0.8.1",
     ],
-    "test": ["mock>=2.0.0", "freezegun==0.3.15", "pylint", "flake8", "coverage", "pytest", "pytest-timeout"],
     "sentry": ["sentry-sdk", "django"],
+    "langchain": ["langchain>=0.2.0"],
 }
 
 setup(
@@ -37,7 +73,16 @@ setup(
     maintainer="PostHog",
     maintainer_email="hey@posthog.com",
     test_suite="posthog.test.all",
-    packages=["posthog", "posthog.test", "posthog.sentry"],
+    packages=[
+        "posthog",
+        "posthog.ai",
+        "posthog.ai.langchain",
+        "posthog.ai.openai",
+        "posthog.ai.anthropic",
+        "posthog.test",
+        "posthog.sentry",
+        "posthog.exception_integrations",
+    ],
     license="MIT License",
     install_requires=install_requires,
     extras_require=extras_require,
@@ -49,16 +94,10 @@ setup(
         "License :: OSI Approved :: MIT License",
         "Operating System :: OS Independent",
         "Programming Language :: Python",
-        "Programming Language :: Python :: 2",
-        "Programming Language :: Python :: 2.6",
-        "Programming Language :: Python :: 2.7",
-        "Programming Language :: Python :: 3",
-        "Programming Language :: Python :: 3.2",
-        "Programming Language :: Python :: 3.3",
-        "Programming Language :: Python :: 3.4",
-        "Programming Language :: Python :: 3.5",
-        "Programming Language :: Python :: 3.6",
-        "Programming Language :: Python :: 3.7",
-        "Programming Language :: Python :: 3.8",
+        "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
+        "Programming Language :: Python :: 3.12",
+        "Programming Language :: Python :: 3.13",
     ],
 )

setup_analytics.py

@@ -27,7 +27,16 @@ setup(
     maintainer="PostHog",
     maintainer_email="hey@posthog.com",
     test_suite="posthoganalytics.test.all",
-    packages=["posthoganalytics", "posthoganalytics.test", "posthoganalytics.sentry"],
+    packages=[
+        "posthoganalytics",
+        "posthoganalytics.ai",
+        "posthoganalytics.ai.langchain",
+        "posthoganalytics.ai.openai",
+        "posthoganalytics.ai.anthropic",
+        "posthoganalytics.test",
+        "posthoganalytics.sentry",
+        "posthoganalytics.exception_integrations",
+    ],
     license="MIT License",
     install_requires=install_requires,
     tests_require=tests_require,
@@ -53,5 +62,8 @@ setup(
         "Programming Language :: Python :: 3.6",
         "Programming Language :: Python :: 3.7",
         "Programming Language :: Python :: 3.8",
+        "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
     ],
 )

simulator.py

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