Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
bcf5b27083 | ||
|
|
0b6ff2e8d3 | ||
|
|
80f0b3e52e | ||
|
|
d1e22188ec | ||
|
|
9b423495ed | ||
|
|
7870ccd3d8 | ||
|
|
190c628c7a | ||
|
|
78ab0ca8b5 | ||
|
|
c5bfc1377a | ||
|
|
6b1c0dc313 | ||
|
|
e51b883e7b | ||
|
|
66101c92bf | ||
|
|
05932b3f13 | ||
|
|
50c13563b2 | ||
|
|
dca4af66ae | ||
|
|
9e1bb8c58a | ||
|
|
fb57de2e12 | ||
|
|
db565bc0fd | ||
|
|
8ae3f2b623 | ||
|
|
39f72a0070 | ||
|
|
ee0305993d | ||
|
|
28c4802d9b | ||
|
|
67a343f242 | ||
|
|
1521621d66 | ||
|
|
39070babfb | ||
|
|
716eab0bc2 | ||
|
|
1c0a61d6b5 | ||
|
|
ffa35fa5cd | ||
|
|
24b7b918f7 | ||
|
|
16cbd10f1b | ||
|
|
b83d544931 | ||
|
|
72c0ed1935 | ||
|
|
5fdd6177ee | ||
|
|
fc1da7d589 | ||
|
|
cba6e86537 | ||
|
|
4e45255207 | ||
|
|
bc37351ab4 | ||
|
|
a5e8b7d7fb | ||
|
|
efb0ccf3c7 | ||
|
|
8554b51a48 | ||
|
|
d0d962a8ba | ||
|
|
e348106094 | ||
|
|
e60d52c199 | ||
|
|
a2c73d0536 | ||
|
|
33ba5d6843 | ||
|
|
3515c40483 | ||
|
|
139258cacb | ||
|
|
f75d924d4c | ||
|
|
617bb53501 |
@@ -18,7 +18,7 @@ jobs:
|
||||
with:
|
||||
python-version: 3.8
|
||||
|
||||
- uses: actions/cache@v1
|
||||
- uses: actions/cache@v3
|
||||
with:
|
||||
path: ~/.cache/pip
|
||||
key: ${{ runner.os }}-pip-${{ hashFiles('setup.py') }}
|
||||
@@ -33,10 +33,10 @@ 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: |
|
||||
@@ -47,14 +47,14 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
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 3.9
|
||||
uses: actions/setup-python@v2
|
||||
with:
|
||||
python-version: 3.7
|
||||
python-version: 3.9
|
||||
|
||||
- name: Install requirements.txt dependencies with pip
|
||||
run: |
|
||||
@@ -62,4 +62,4 @@ jobs:
|
||||
|
||||
- name: Run posthog tests
|
||||
run: |
|
||||
python setup.py test
|
||||
pytest --verbose --timeout=30
|
||||
|
||||
+3
-1
@@ -14,4 +14,6 @@ pylint.out
|
||||
posthog-analytics
|
||||
.idea
|
||||
.python-version
|
||||
.coverage
|
||||
.coverage
|
||||
pyrightconfig.json
|
||||
.env
|
||||
|
||||
+166
-5
@@ -1,3 +1,161 @@
|
||||
## 3.8.5 - 2025-01-22
|
||||
|
||||
1. Add tracing event to langchain callbacks.
|
||||
|
||||
## 3.8.4 - 2025-01-17
|
||||
|
||||
1. Add Anthropic support for LLM Observability.
|
||||
2. Update LLM Observability to use output_choices.
|
||||
|
||||
## 3.8.3 - 2025-01-14
|
||||
|
||||
1. Fix setuptools to include the `posthog.ai.openai` and `posthog.ai.langchain` packages for the `posthoganalytics` package.
|
||||
|
||||
## 3.8.2 - 2025-01-14
|
||||
|
||||
1. Fix setuptools to include the `posthog.ai.openai` and `posthog.ai.langchain` packages.
|
||||
|
||||
## 3.8.1 - 2025-01-14
|
||||
|
||||
1. Add LLM Observability with support for OpenAI and Langchain callbacks.
|
||||
|
||||
## 3.7.5 - 2025-01-03
|
||||
|
||||
1. Add `distinct_id` to group_identify
|
||||
|
||||
## 3.7.4 - 2024-11-25
|
||||
|
||||
1. Fix bug where this SDK incorrectly sent feature flag events with null values when calling `get_feature_flag_payload`.
|
||||
|
||||
## 3.7.3 - 2024-11-25
|
||||
|
||||
1. Use personless mode when sending an exception without a provided `distinct_id`.
|
||||
|
||||
## 3.7.2 - 2024-11-19
|
||||
|
||||
1. Add `type` property to exception stacks.
|
||||
|
||||
## 3.7.1 - 2024-10-24
|
||||
|
||||
1. Add `platform` property to each frame of exception stacks.
|
||||
|
||||
## 3.7.0 - 2024-10-03
|
||||
|
||||
1. Adds a new `super_properties` parameter on the client that are appended to every /capture call.
|
||||
|
||||
## 3.6.7 - 2024-09-24
|
||||
|
||||
1. Remove deprecated datetime.utcnow() in favour of datetime.now(tz=tzutc())
|
||||
|
||||
## 3.6.6 - 2024-09-16
|
||||
|
||||
1. Fix manual capture support for in app frames
|
||||
|
||||
## 3.6.5 - 2024-09-10
|
||||
|
||||
1. Fix django integration support for manual exception capture.
|
||||
|
||||
## 3.6.4 - 2024-09-05
|
||||
|
||||
1. Add manual exception capture.
|
||||
|
||||
## 3.6.3 - 2024-09-03
|
||||
|
||||
1. Make sure setup.py for posthoganalytics package also discovers the new exception integration package.
|
||||
|
||||
## 3.6.2 - 2024-09-03
|
||||
|
||||
1. Make sure setup.py discovers the new exception integration package.
|
||||
|
||||
## 3.6.1 - 2024-09-03
|
||||
|
||||
1. Adds django integration to exception autocapture in alpha state. This feature is not yet stable and may change in future versions.
|
||||
|
||||
## 3.6.0 - 2024-08-28
|
||||
|
||||
1. Adds exception autocapture in alpha state. This feature is not yet stable and may change in future versions.
|
||||
|
||||
## 3.5.2 - 2024-08-21
|
||||
|
||||
1. Guard for None values in local evaluation
|
||||
|
||||
## 3.5.1 - 2024-08-13
|
||||
|
||||
1. Remove "-api" suffix from ingestion hostnames
|
||||
|
||||
## 3.5.0 - 2024-02-29
|
||||
|
||||
1. - Adds a new `feature_flags_request_timeout_seconds` timeout parameter for feature flags which defaults to 3 seconds, updated from the default 10s for all other API calls.
|
||||
|
||||
## 3.4.2 - 2024-02-20
|
||||
|
||||
1. Add `historical_migration` option for bulk migration to PostHog Cloud.
|
||||
|
||||
## 3.4.1 - 2024-02-09
|
||||
|
||||
1. Use new hosts for event capture as well
|
||||
|
||||
## 3.4.0 - 2024-02-05
|
||||
|
||||
1. Point given hosts to new ingestion hosts
|
||||
|
||||
## 3.3.4 - 2024-01-30
|
||||
|
||||
1. Update type hints for module variables to work with newer versions of mypy
|
||||
|
||||
## 3.3.3 - 2024-01-26
|
||||
|
||||
1. Remove new relative date operators, combine into regular date operators
|
||||
|
||||
## 3.3.2 - 2024-01-19
|
||||
|
||||
1. Return success/failure with all capture calls from module functions
|
||||
|
||||
## 3.3.1 - 2024-01-10
|
||||
|
||||
1. Make sure we don't override any existing feature flag properties when adding locally evaluated feature flag properties.
|
||||
|
||||
## 3.3.0 - 2024-01-09
|
||||
|
||||
1. When local evaluation is enabled, we automatically add flag information to all events sent to PostHog, whenever possible. This makes it easier to use these events in experiments.
|
||||
|
||||
## 3.2.0 - 2024-01-09
|
||||
|
||||
1. Numeric property handling for feature flags now does the expected: When passed in a number, we do a numeric comparison. When passed in a string, we do a string comparison. Previously, we always did a string comparison.
|
||||
2. Add support for relative date operators for local evaluation.
|
||||
|
||||
## 3.1.0 - 2023-12-04
|
||||
|
||||
1. Increase maximum event size and batch size
|
||||
|
||||
## 3.0.2 - 2023-08-17
|
||||
|
||||
1. Returns the current flag property with $feature_flag_called events, to make it easier to use in experiments
|
||||
|
||||
## 3.0.1 - 2023-04-21
|
||||
|
||||
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:
|
||||
|
||||
All events by default now send the `$geoip_disable` property to disable geoip lookup in app. This is because usually we don't
|
||||
want to update person properties to take the server's location.
|
||||
|
||||
The same now happens for feature flag requests, where we discard the IP address of the server for matching on geoip properties like city, country, continent.
|
||||
|
||||
To restore previous behaviour, you can set the default to False like so:
|
||||
|
||||
```python
|
||||
posthog.disable_geoip = False
|
||||
|
||||
# // and if using client instantiation:
|
||||
posthog = Posthog('api_key', disable_geoip=False)
|
||||
|
||||
```
|
||||
|
||||
## 2.5.0 - 2023-04-10
|
||||
|
||||
1. Add option for instantiating separate client object
|
||||
@@ -19,7 +177,6 @@
|
||||
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
|
||||
@@ -35,6 +192,7 @@ Changes:
|
||||
Changes:
|
||||
|
||||
1. Fixes issues with date comparison.
|
||||
|
||||
## 2.1.1 - 2022-09-14
|
||||
|
||||
Changes:
|
||||
@@ -47,6 +205,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
|
||||
@@ -58,8 +217,8 @@ Changes:
|
||||
Breaking changes:
|
||||
|
||||
1. The minimum version requirement for PostHog servers is now 1.38. If you're using PostHog Cloud, you satisfy this requirement automatically.
|
||||
2. Feature flag defaults apply only when there's an error fetching feature flag results. Earlier, if the default was set to `True`, even if a flag resolved to `False`, the default would override this.
|
||||
**Note: These are removed in 2.0.2**
|
||||
2. Feature flag defaults apply only when there's an error fetching feature flag results. Earlier, if the default was set to `True`, even if a flag resolved to `False`, the default would override this.
|
||||
**Note: These are removed in 2.0.2**
|
||||
3. Feature flag remote evaluation doesn't require a personal API key.
|
||||
|
||||
New Changes:
|
||||
@@ -67,18 +226,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
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
@PostHog/team-feature-flags
|
||||
@@ -1,4 +1,4 @@
|
||||
Copyright (c) 2020 PostHog (part of Hiberly Inc)
|
||||
Copyright (c) 2023 PostHog (part of Hiberly Inc)
|
||||
|
||||
Copyright (c) 2013 Segment Inc. friends@segment.com
|
||||
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
# PostHog Python library example
|
||||
|
||||
# Import the library
|
||||
import time
|
||||
# import time
|
||||
|
||||
import posthog
|
||||
|
||||
|
||||
@@ -0,0 +1,196 @@
|
||||
import os
|
||||
import uuid
|
||||
|
||||
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.personal_api_key = os.getenv("POSTHOG_PERSONAL_API_KEY", "your-personal-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()
|
||||
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
|
||||
|
||||
|
||||
# HOW TO RUN:
|
||||
# comment out one of these to run the other
|
||||
|
||||
if __name__ == "__main__":
|
||||
main_sync()
|
||||
|
||||
# asyncio.run(main_async())
|
||||
+114
-19
@@ -1,24 +1,35 @@
|
||||
import datetime # noqa: F401
|
||||
from typing import Callable, Dict, Optional # noqa: F401
|
||||
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
|
||||
|
||||
"""Settings."""
|
||||
api_key = None # type: str
|
||||
host = None # type: str
|
||||
on_error = None # type: Callable
|
||||
api_key = None # type: Optional[str]
|
||||
host = None # type: Optional[str]
|
||||
on_error = None # type: Optional[Callable]
|
||||
debug = False # type: bool
|
||||
send = True # type: bool
|
||||
sync_mode = False # type: bool
|
||||
disabled = False # type: bool
|
||||
personal_api_key = None # type: str
|
||||
project_api_key = None # type: str
|
||||
personal_api_key = None # type: Optional[str]
|
||||
project_api_key = None # type: Optional[str]
|
||||
poll_interval = 30 # type: int
|
||||
disable_geoip = True # type: bool
|
||||
feature_flags_request_timeout_seconds = 3 # type: int
|
||||
super_properties = None # type: Optional[Dict]
|
||||
# Currently alpha, use at your own risk
|
||||
enable_exception_autocapture = False # type: bool
|
||||
exception_autocapture_integrations = [] # type: List[Integrations]
|
||||
# Used to determine in app paths for exception autocapture. Defaults to the current working directory
|
||||
project_root = None # type: Optional[str]
|
||||
# Used for our AI observability feature to not capture any prompt or output just usage + metadata
|
||||
privacy_mode = False # type: bool
|
||||
|
||||
default_client = None
|
||||
default_client = None # type: Optional[Client]
|
||||
|
||||
|
||||
def capture(
|
||||
@@ -30,8 +41,9 @@ def capture(
|
||||
uuid=None, # type: Optional[str]
|
||||
groups=None, # type: Optional[Dict]
|
||||
send_feature_flags=False,
|
||||
disable_geoip=None, # type: Optional[bool]
|
||||
):
|
||||
# type: (...) -> None
|
||||
# type: (...) -> Tuple[bool, dict]
|
||||
"""
|
||||
Capture allows you to capture anything a user does within your system, which you can later use in PostHog to find patterns in usage, work out which features to improve or where people are giving up.
|
||||
|
||||
@@ -52,7 +64,7 @@ def capture(
|
||||
posthog.capture('distinct id', 'purchase', groups={'company': 'id:5'})
|
||||
```
|
||||
"""
|
||||
_proxy(
|
||||
return _proxy(
|
||||
"capture",
|
||||
distinct_id=distinct_id,
|
||||
event=event,
|
||||
@@ -62,6 +74,7 @@ def capture(
|
||||
uuid=uuid,
|
||||
groups=groups,
|
||||
send_feature_flags=send_feature_flags,
|
||||
disable_geoip=disable_geoip,
|
||||
)
|
||||
|
||||
|
||||
@@ -71,8 +84,9 @@ def identify(
|
||||
context=None, # type: Optional[Dict]
|
||||
timestamp=None, # type: Optional[datetime.datetime]
|
||||
uuid=None, # type: Optional[str]
|
||||
disable_geoip=None, # type: Optional[bool]
|
||||
):
|
||||
# type: (...) -> None
|
||||
# type: (...) -> Tuple[bool, dict]
|
||||
"""
|
||||
Identify lets you add metadata on your users so you can more easily identify who they are in PostHog, and even do things like segment users by these properties.
|
||||
|
||||
@@ -88,13 +102,14 @@ def identify(
|
||||
})
|
||||
```
|
||||
"""
|
||||
_proxy(
|
||||
return _proxy(
|
||||
"identify",
|
||||
distinct_id=distinct_id,
|
||||
properties=properties,
|
||||
context=context,
|
||||
timestamp=timestamp,
|
||||
uuid=uuid,
|
||||
disable_geoip=disable_geoip,
|
||||
)
|
||||
|
||||
|
||||
@@ -104,8 +119,9 @@ def set(
|
||||
context=None, # type: Optional[Dict]
|
||||
timestamp=None, # type: Optional[datetime.datetime]
|
||||
uuid=None, # type: Optional[str]
|
||||
disable_geoip=None, # type: Optional[bool]
|
||||
):
|
||||
# type: (...) -> None
|
||||
# type: (...) -> Tuple[bool, dict]
|
||||
"""
|
||||
Set properties on a user record.
|
||||
This will overwrite previous people property values, just like `identify`.
|
||||
@@ -121,13 +137,14 @@ def set(
|
||||
})
|
||||
```
|
||||
"""
|
||||
_proxy(
|
||||
return _proxy(
|
||||
"set",
|
||||
distinct_id=distinct_id,
|
||||
properties=properties,
|
||||
context=context,
|
||||
timestamp=timestamp,
|
||||
uuid=uuid,
|
||||
disable_geoip=disable_geoip,
|
||||
)
|
||||
|
||||
|
||||
@@ -137,8 +154,9 @@ def set_once(
|
||||
context=None, # type: Optional[Dict]
|
||||
timestamp=None, # type: Optional[datetime.datetime]
|
||||
uuid=None, # type: Optional[str]
|
||||
disable_geoip=None, # type: Optional[bool]
|
||||
):
|
||||
# type: (...) -> None
|
||||
# type: (...) -> Tuple[bool, dict]
|
||||
"""
|
||||
Set properties on a user record, only if they do not yet exist.
|
||||
This will not overwrite previous people property values, unlike `identify`.
|
||||
@@ -154,13 +172,14 @@ def set_once(
|
||||
})
|
||||
```
|
||||
"""
|
||||
_proxy(
|
||||
return _proxy(
|
||||
"set_once",
|
||||
distinct_id=distinct_id,
|
||||
properties=properties,
|
||||
context=context,
|
||||
timestamp=timestamp,
|
||||
uuid=uuid,
|
||||
disable_geoip=disable_geoip,
|
||||
)
|
||||
|
||||
|
||||
@@ -171,8 +190,9 @@ def group_identify(
|
||||
context=None, # type: Optional[Dict]
|
||||
timestamp=None, # type: Optional[datetime.datetime]
|
||||
uuid=None, # type: Optional[str]
|
||||
disable_geoip=None, # type: Optional[bool]
|
||||
):
|
||||
# type: (...) -> None
|
||||
# type: (...) -> Tuple[bool, dict]
|
||||
"""
|
||||
Set properties on a group
|
||||
|
||||
@@ -188,7 +208,7 @@ def group_identify(
|
||||
})
|
||||
```
|
||||
"""
|
||||
_proxy(
|
||||
return _proxy(
|
||||
"group_identify",
|
||||
group_type=group_type,
|
||||
group_key=group_key,
|
||||
@@ -196,6 +216,7 @@ def group_identify(
|
||||
context=context,
|
||||
timestamp=timestamp,
|
||||
uuid=uuid,
|
||||
disable_geoip=disable_geoip,
|
||||
)
|
||||
|
||||
|
||||
@@ -205,8 +226,9 @@ def alias(
|
||||
context=None, # type: Optional[Dict]
|
||||
timestamp=None, # type: Optional[datetime.datetime]
|
||||
uuid=None, # type: Optional[str]
|
||||
disable_geoip=None, # type: Optional[bool]
|
||||
):
|
||||
# type: (...) -> None
|
||||
# type: (...) -> Tuple[bool, dict]
|
||||
"""
|
||||
To marry up whatever a user does before they sign up or log in with what they do after you need to make an alias call. This will allow you to answer questions like "Which marketing channels leads to users churning after a month?" or "What do users do on our website before signing up?"
|
||||
|
||||
@@ -223,13 +245,58 @@ def alias(
|
||||
posthog.alias('anonymous session id', 'distinct id')
|
||||
```
|
||||
"""
|
||||
_proxy(
|
||||
return _proxy(
|
||||
"alias",
|
||||
previous_id=previous_id,
|
||||
distinct_id=distinct_id,
|
||||
context=context,
|
||||
timestamp=timestamp,
|
||||
uuid=uuid,
|
||||
disable_geoip=disable_geoip,
|
||||
)
|
||||
|
||||
|
||||
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')
|
||||
|
||||
```
|
||||
"""
|
||||
return _proxy(
|
||||
"capture_exception",
|
||||
exception=exception,
|
||||
distinct_id=distinct_id,
|
||||
properties=properties,
|
||||
context=context,
|
||||
timestamp=timestamp,
|
||||
uuid=uuid,
|
||||
groups=groups,
|
||||
)
|
||||
|
||||
|
||||
@@ -241,6 +308,7 @@ def feature_enabled(
|
||||
group_properties={}, # type: dict
|
||||
only_evaluate_locally=False, # type: bool
|
||||
send_feature_flag_events=True, # type: bool
|
||||
disable_geoip=None, # type: Optional[bool]
|
||||
):
|
||||
# type: (...) -> bool
|
||||
"""
|
||||
@@ -265,6 +333,7 @@ def feature_enabled(
|
||||
group_properties=group_properties,
|
||||
only_evaluate_locally=only_evaluate_locally,
|
||||
send_feature_flag_events=send_feature_flag_events,
|
||||
disable_geoip=disable_geoip,
|
||||
)
|
||||
|
||||
|
||||
@@ -276,6 +345,7 @@ def get_feature_flag(
|
||||
group_properties={}, # type: dict
|
||||
only_evaluate_locally=False, # type: bool
|
||||
send_feature_flag_events=True, # type: bool
|
||||
disable_geoip=None, # type: Optional[bool]
|
||||
):
|
||||
"""
|
||||
Get feature flag variant for users. Used with experiments.
|
||||
@@ -308,6 +378,7 @@ def get_feature_flag(
|
||||
group_properties=group_properties,
|
||||
only_evaluate_locally=only_evaluate_locally,
|
||||
send_feature_flag_events=send_feature_flag_events,
|
||||
disable_geoip=disable_geoip,
|
||||
)
|
||||
|
||||
|
||||
@@ -317,6 +388,7 @@ def get_all_flags(
|
||||
person_properties={}, # type: dict
|
||||
group_properties={}, # type: dict
|
||||
only_evaluate_locally=False, # type: bool
|
||||
disable_geoip=None, # type: Optional[bool]
|
||||
):
|
||||
"""
|
||||
Get all flags for a given user.
|
||||
@@ -334,6 +406,7 @@ def get_all_flags(
|
||||
person_properties=person_properties,
|
||||
group_properties=group_properties,
|
||||
only_evaluate_locally=only_evaluate_locally,
|
||||
disable_geoip=disable_geoip,
|
||||
)
|
||||
|
||||
|
||||
@@ -346,6 +419,7 @@ def get_feature_flag_payload(
|
||||
group_properties={},
|
||||
only_evaluate_locally=False,
|
||||
send_feature_flag_events=True,
|
||||
disable_geoip=None, # type: Optional[bool]
|
||||
):
|
||||
return _proxy(
|
||||
"get_feature_flag_payload",
|
||||
@@ -357,6 +431,7 @@ def get_feature_flag_payload(
|
||||
group_properties=group_properties,
|
||||
only_evaluate_locally=only_evaluate_locally,
|
||||
send_feature_flag_events=send_feature_flag_events,
|
||||
disable_geoip=disable_geoip,
|
||||
)
|
||||
|
||||
|
||||
@@ -366,6 +441,7 @@ def get_all_flags_and_payloads(
|
||||
person_properties={},
|
||||
group_properties={},
|
||||
only_evaluate_locally=False,
|
||||
disable_geoip=None, # type: Optional[bool]
|
||||
):
|
||||
return _proxy(
|
||||
"get_all_flags_and_payloads",
|
||||
@@ -374,9 +450,20 @@ def get_all_flags_and_payloads(
|
||||
person_properties=person_properties,
|
||||
group_properties=group_properties,
|
||||
only_evaluate_locally=only_evaluate_locally,
|
||||
disable_geoip=disable_geoip,
|
||||
)
|
||||
|
||||
|
||||
def feature_flag_definitions():
|
||||
"""Returns loaded feature flags, if any. Helpful for debugging what flag information you have loaded."""
|
||||
return _proxy("feature_flag_definitions")
|
||||
|
||||
|
||||
def load_feature_flags():
|
||||
"""Load feature flag definitions from PostHog."""
|
||||
return _proxy("load_feature_flags")
|
||||
|
||||
|
||||
def page(*args, **kwargs):
|
||||
"""Send a page call."""
|
||||
_proxy("page", *args, **kwargs)
|
||||
@@ -418,6 +505,14 @@ def _proxy(method, *args, **kwargs):
|
||||
project_api_key=project_api_key,
|
||||
poll_interval=poll_interval,
|
||||
disabled=disabled,
|
||||
disable_geoip=disable_geoip,
|
||||
feature_flags_request_timeout_seconds=feature_flags_request_timeout_seconds,
|
||||
super_properties=super_properties,
|
||||
# TODO: Currently this monitoring begins only when the Client is initialised (which happens when you do something with the SDK)
|
||||
# This kind of initialisation is very annoying for exception capture. We need to figure out a way around this,
|
||||
# or deprecate this proxy option fully (it's already in the process of deprecation, no new clients should be using this method since like 5-6 months)
|
||||
enable_exception_autocapture=enable_exception_autocapture,
|
||||
exception_autocapture_integrations=exception_autocapture_integrations,
|
||||
)
|
||||
|
||||
# always set incase user changes it
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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)
|
||||
@@ -0,0 +1,3 @@
|
||||
from .callbacks import CallbackHandler
|
||||
|
||||
__all__ = ["CallbackHandler"]
|
||||
@@ -0,0 +1,597 @@
|
||||
try:
|
||||
import langchain # noqa: F401
|
||||
except ImportError:
|
||||
raise ModuleNotFoundError("Please install LangChain to use this feature: 'pip install langchain'")
|
||||
|
||||
import logging
|
||||
import time
|
||||
import uuid
|
||||
from typing import (
|
||||
Any,
|
||||
Dict,
|
||||
List,
|
||||
Optional,
|
||||
Tuple,
|
||||
TypedDict,
|
||||
Union,
|
||||
cast,
|
||||
)
|
||||
from uuid import UUID
|
||||
|
||||
from langchain.callbacks.base import BaseCallbackHandler
|
||||
from langchain.schema.agent import AgentAction, AgentFinish
|
||||
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")
|
||||
|
||||
|
||||
class RunMetadata(TypedDict, total=False):
|
||||
messages: Union[List[Dict[str, Any]], List[str]]
|
||||
provider: str
|
||||
model: str
|
||||
model_params: Dict[str, Any]
|
||||
base_url: str
|
||||
start_time: float
|
||||
end_time: float
|
||||
|
||||
|
||||
RunStorage = 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: RunStorage
|
||||
"""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.
|
||||
"""
|
||||
self._client = client or default_client
|
||||
self._distinct_id = distinct_id
|
||||
self._trace_id = trace_id
|
||||
self._trace_name = None
|
||||
self._trace_input = None
|
||||
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)
|
||||
if parent_run_id is None and self._trace_name is None:
|
||||
self._trace_name = self._get_langchain_run_name(serialized, **kwargs)
|
||||
self._trace_input = inputs
|
||||
|
||||
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_run_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_run_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_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)
|
||||
|
||||
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)
|
||||
|
||||
def on_tool_error(
|
||||
self,
|
||||
error: Union[Exception, KeyboardInterrupt],
|
||||
*,
|
||||
run_id: UUID,
|
||||
parent_run_id: Optional[UUID] = None,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
self._log_debug_event("on_tool_error", run_id, parent_run_id, error=error)
|
||||
|
||||
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_parent_of_run(run_id)
|
||||
|
||||
if parent_run_id is None:
|
||||
self._capture_trace(run_id, outputs=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_parent_of_run(run_id)
|
||||
|
||||
if parent_run_id is None:
|
||||
self._capture_trace(run_id, outputs=None)
|
||||
|
||||
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)
|
||||
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
|
||||
|
||||
latency = run.get("end_time", 0) - run.get("start_time", 0)
|
||||
input_tokens, output_tokens = _parse_usage(response)
|
||||
|
||||
generation_result = response.generations[-1]
|
||||
if isinstance(generation_result[-1], ChatGeneration):
|
||||
output = [
|
||||
_convert_message_to_dict(cast(ChatGeneration, generation).message) for generation in generation_result
|
||||
]
|
||||
else:
|
||||
output = [_extract_raw_esponse(generation) for generation in generation_result]
|
||||
|
||||
event_properties = {
|
||||
"$ai_provider": run.get("provider"),
|
||||
"$ai_model": run.get("model"),
|
||||
"$ai_model_parameters": run.get("model_params"),
|
||||
"$ai_input": with_privacy_mode(self._client, self._privacy_mode, run.get("messages")),
|
||||
"$ai_output_choices": with_privacy_mode(self._client, self._privacy_mode, output),
|
||||
"$ai_http_status": 200,
|
||||
"$ai_input_tokens": input_tokens,
|
||||
"$ai_output_tokens": output_tokens,
|
||||
"$ai_latency": latency,
|
||||
"$ai_trace_id": trace_id,
|
||||
"$ai_base_url": run.get("base_url"),
|
||||
**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 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)
|
||||
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
|
||||
|
||||
latency = run.get("end_time", 0) - run.get("start_time", 0)
|
||||
event_properties = {
|
||||
"$ai_provider": run.get("provider"),
|
||||
"$ai_model": run.get("model"),
|
||||
"$ai_model_parameters": run.get("model_params"),
|
||||
"$ai_input": with_privacy_mode(self._client, self._privacy_mode, run.get("messages")),
|
||||
"$ai_http_status": _get_http_status(error),
|
||||
"$ai_latency": latency,
|
||||
"$ai_trace_id": trace_id,
|
||||
"$ai_base_url": run.get("base_url"),
|
||||
**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 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)
|
||||
|
||||
def on_retriever_error(
|
||||
self,
|
||||
error: Union[Exception, KeyboardInterrupt],
|
||||
*,
|
||||
run_id: UUID,
|
||||
parent_run_id: Optional[UUID] = None,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
"""Run when Retriever errors."""
|
||||
self._log_debug_event("on_retriever_error", run_id, parent_run_id, error=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)
|
||||
|
||||
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)
|
||||
|
||||
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_run_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: RunMetadata = {
|
||||
"messages": messages,
|
||||
"start_time": time.time(),
|
||||
}
|
||||
if isinstance(invocation_params, dict):
|
||||
run["model_params"] = get_model_params(invocation_params)
|
||||
if isinstance(metadata, dict):
|
||||
if model := metadata.get("ls_model_name"):
|
||||
run["model"] = model
|
||||
if provider := metadata.get("ls_provider"):
|
||||
run["provider"] = provider
|
||||
try:
|
||||
base_url = serialized["kwargs"]["openai_api_base"]
|
||||
if base_url is not None:
|
||||
run["base_url"] = base_url
|
||||
except KeyError:
|
||||
pass
|
||||
self._runs[run_id] = run
|
||||
|
||||
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:
|
||||
trace_id = uuid.uuid4()
|
||||
return trace_id
|
||||
|
||||
def _get_langchain_run_name(self, serialized: Optional[Dict[str, Any]], **kwargs: Any) -> 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"]
|
||||
|
||||
try:
|
||||
return serialized["name"]
|
||||
except (KeyError, TypeError):
|
||||
pass
|
||||
|
||||
try:
|
||||
return serialized["id"][-1]
|
||||
except (KeyError, TypeError):
|
||||
pass
|
||||
|
||||
def _capture_trace(self, run_id: UUID, *, outputs: Optional[Dict[str, Any]]):
|
||||
trace_id = self._get_trace_id(run_id)
|
||||
event_properties = {
|
||||
"$ai_trace_name": self._trace_name,
|
||||
"$ai_trace_id": trace_id,
|
||||
"$ai_input_state": with_privacy_mode(self._client, self._privacy_mode, self._trace_input),
|
||||
**self._properties,
|
||||
}
|
||||
if 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 trace_id,
|
||||
event="$ai_trace",
|
||||
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"),
|
||||
# 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:
|
||||
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
|
||||
@@ -0,0 +1,4 @@
|
||||
from .openai import OpenAI
|
||||
from .openai_async import AsyncOpenAI
|
||||
|
||||
__all__ = ["OpenAI", "AsyncOpenAI"]
|
||||
@@ -0,0 +1,251 @@
|
||||
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)
|
||||
|
||||
|
||||
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
|
||||
@@ -0,0 +1,250 @@
|
||||
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)
|
||||
|
||||
|
||||
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,
|
||||
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
|
||||
@@ -0,0 +1,245 @@
|
||||
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] = {}
|
||||
|
||||
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
|
||||
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 {}),
|
||||
}
|
||||
|
||||
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] = {}
|
||||
|
||||
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
|
||||
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 {}),
|
||||
}
|
||||
|
||||
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
|
||||
+322
-56
@@ -1,17 +1,21 @@
|
||||
import atexit
|
||||
import logging
|
||||
import numbers
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime, timedelta
|
||||
from uuid import UUID
|
||||
from uuid import UUID, uuid4
|
||||
|
||||
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, get
|
||||
from posthog.utils import SizeLimitedDict, clean, guess_timezone
|
||||
from posthog.request import DEFAULT_HOST, APIError, batch_post, decide, determine_server_host, get
|
||||
from posthog.utils import SizeLimitedDict, clean, guess_timezone, remove_trailing_slash
|
||||
from posthog.version import VERSION
|
||||
|
||||
try:
|
||||
@@ -48,6 +52,14 @@ class Client(object):
|
||||
personal_api_key=None,
|
||||
project_api_key=None,
|
||||
disabled=False,
|
||||
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)
|
||||
|
||||
@@ -60,7 +72,9 @@ class Client(object):
|
||||
self.debug = debug
|
||||
self.send = send
|
||||
self.sync_mode = sync_mode
|
||||
self.host = host
|
||||
# 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
|
||||
self.feature_flags = None
|
||||
@@ -68,9 +82,25 @@ class Client(object):
|
||||
self.group_type_mapping = None
|
||||
self.cohorts = None
|
||||
self.poll_interval = poll_interval
|
||||
self.feature_flags_request_timeout_seconds = feature_flags_request_timeout_seconds
|
||||
self.poller = None
|
||||
self.distinct_ids_feature_flags_reported = SizeLimitedDict(MAX_DICT_SIZE, set)
|
||||
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
|
||||
@@ -82,6 +112,9 @@ class Client(object):
|
||||
else:
|
||||
self.log.setLevel(logging.WARNING)
|
||||
|
||||
if self.enable_exception_autocapture:
|
||||
self.exception_capture = ExceptionCapture(self, integrations=self.exception_autocapture_integrations)
|
||||
|
||||
if sync_mode:
|
||||
self.consumers = None
|
||||
else:
|
||||
@@ -98,13 +131,14 @@ class Client(object):
|
||||
consumer = Consumer(
|
||||
self.queue,
|
||||
self.api_key,
|
||||
host=host,
|
||||
host=self.host,
|
||||
on_error=on_error,
|
||||
flush_at=flush_at,
|
||||
flush_interval=flush_interval,
|
||||
gzip=gzip,
|
||||
retries=max_retries,
|
||||
timeout=timeout,
|
||||
historical_migration=historical_migration,
|
||||
)
|
||||
self.consumers.append(consumer)
|
||||
|
||||
@@ -112,7 +146,7 @@ class Client(object):
|
||||
if send:
|
||||
consumer.start()
|
||||
|
||||
def identify(self, distinct_id=None, properties=None, context=None, timestamp=None, uuid=None):
|
||||
def identify(self, distinct_id=None, properties=None, context=None, timestamp=None, uuid=None, disable_geoip=None):
|
||||
properties = properties or {}
|
||||
context = context or {}
|
||||
require("distinct_id", distinct_id, ID_TYPES)
|
||||
@@ -127,25 +161,35 @@ class Client(object):
|
||||
"uuid": uuid,
|
||||
}
|
||||
|
||||
return self._enqueue(msg)
|
||||
return self._enqueue(msg, disable_geoip)
|
||||
|
||||
def get_feature_variants(self, distinct_id, groups=None, person_properties=None, group_properties=None):
|
||||
resp_data = self.get_decide(distinct_id, groups, person_properties, group_properties)
|
||||
def get_feature_variants(
|
||||
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 resp_data["featureFlags"]
|
||||
|
||||
def _get_active_feature_variants(self, distinct_id, groups=None, person_properties=None, group_properties=None):
|
||||
feature_variants = self.get_feature_variants(distinct_id, groups, person_properties, group_properties)
|
||||
return {
|
||||
k: v for (k, v) in feature_variants.items() if v is not False
|
||||
} # explicitly test for false to account for values that may seem falsy (ex: 0)
|
||||
|
||||
def get_feature_payloads(self, distinct_id, groups=None, person_properties=None, group_properties=None):
|
||||
resp_data = self.get_decide(distinct_id, groups, person_properties, group_properties)
|
||||
def get_feature_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 resp_data["featureFlagPayloads"]
|
||||
|
||||
def get_decide(self, distinct_id, groups=None, person_properties=None, group_properties=None):
|
||||
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)
|
||||
|
||||
if disable_geoip is None:
|
||||
disable_geoip = self.disable_geoip
|
||||
|
||||
if groups:
|
||||
require("groups", groups, dict)
|
||||
else:
|
||||
@@ -156,8 +200,9 @@ class Client(object):
|
||||
"groups": groups,
|
||||
"person_properties": person_properties,
|
||||
"group_properties": group_properties,
|
||||
"disable_geoip": disable_geoip,
|
||||
}
|
||||
resp_data = decide(self.api_key, self.host, timeout=10, **request_data)
|
||||
resp_data = decide(self.api_key, self.host, timeout=self.feature_flags_request_timeout_seconds, **request_data)
|
||||
|
||||
return resp_data
|
||||
|
||||
@@ -171,6 +216,7 @@ class Client(object):
|
||||
uuid=None,
|
||||
groups=None,
|
||||
send_feature_flags=False,
|
||||
disable_geoip=None,
|
||||
):
|
||||
properties = properties or {}
|
||||
context = context or {}
|
||||
@@ -191,19 +237,33 @@ class Client(object):
|
||||
require("groups", groups, dict)
|
||||
msg["properties"]["$groups"] = groups
|
||||
|
||||
extra_properties = {}
|
||||
feature_variants = {}
|
||||
if send_feature_flags:
|
||||
try:
|
||||
feature_variants = self._get_active_feature_variants(distinct_id, groups)
|
||||
feature_variants = self.get_feature_variants(distinct_id, groups, disable_geoip=disable_geoip)
|
||||
except Exception as e:
|
||||
self.log.exception(f"[FEATURE FLAGS] Unable to get feature variants: {e}")
|
||||
else:
|
||||
for feature, variant in feature_variants.items():
|
||||
msg["properties"]["$feature/{}".format(feature)] = variant
|
||||
msg["properties"]["$active_feature_flags"] = list(feature_variants.keys())
|
||||
|
||||
return self._enqueue(msg)
|
||||
elif self.feature_flags:
|
||||
# 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
|
||||
)
|
||||
|
||||
def set(self, distinct_id=None, properties=None, context=None, timestamp=None, uuid=None):
|
||||
for feature, variant in feature_variants.items():
|
||||
extra_properties[f"$feature/{feature}"] = variant
|
||||
|
||||
active_feature_flags = [key for (key, value) in feature_variants.items() if value is not False]
|
||||
if active_feature_flags:
|
||||
extra_properties["$active_feature_flags"] = active_feature_flags
|
||||
|
||||
if extra_properties:
|
||||
msg["properties"] = {**extra_properties, **msg["properties"]}
|
||||
|
||||
return self._enqueue(msg, disable_geoip)
|
||||
|
||||
def set(self, distinct_id=None, properties=None, context=None, timestamp=None, uuid=None, disable_geoip=None):
|
||||
properties = properties or {}
|
||||
context = context or {}
|
||||
require("distinct_id", distinct_id, ID_TYPES)
|
||||
@@ -218,9 +278,9 @@ class Client(object):
|
||||
"uuid": uuid,
|
||||
}
|
||||
|
||||
return self._enqueue(msg)
|
||||
return self._enqueue(msg, disable_geoip)
|
||||
|
||||
def set_once(self, distinct_id=None, properties=None, context=None, timestamp=None, uuid=None):
|
||||
def set_once(self, distinct_id=None, properties=None, context=None, timestamp=None, uuid=None, disable_geoip=None):
|
||||
properties = properties or {}
|
||||
context = context or {}
|
||||
require("distinct_id", distinct_id, ID_TYPES)
|
||||
@@ -235,15 +295,30 @@ class Client(object):
|
||||
"uuid": uuid,
|
||||
}
|
||||
|
||||
return self._enqueue(msg)
|
||||
return self._enqueue(msg, disable_geoip)
|
||||
|
||||
def group_identify(self, group_type=None, group_key=None, properties=None, context=None, timestamp=None, uuid=None):
|
||||
def group_identify(
|
||||
self,
|
||||
group_type=None,
|
||||
group_key=None,
|
||||
properties=None,
|
||||
context=None,
|
||||
timestamp=None,
|
||||
uuid=None,
|
||||
disable_geoip=None,
|
||||
distinct_id=None,
|
||||
):
|
||||
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": {
|
||||
@@ -251,15 +326,15 @@ 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)
|
||||
return self._enqueue(msg, disable_geoip)
|
||||
|
||||
def alias(self, previous_id=None, distinct_id=None, context=None, timestamp=None, uuid=None):
|
||||
def alias(self, previous_id=None, distinct_id=None, context=None, timestamp=None, uuid=None, disable_geoip=None):
|
||||
context = context or {}
|
||||
|
||||
require("previous_id", previous_id, ID_TYPES)
|
||||
@@ -276,9 +351,11 @@ class Client(object):
|
||||
"distinct_id": previous_id,
|
||||
}
|
||||
|
||||
return self._enqueue(msg)
|
||||
return self._enqueue(msg, disable_geoip)
|
||||
|
||||
def page(self, distinct_id=None, url=None, properties=None, context=None, timestamp=None, uuid=None):
|
||||
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 {}
|
||||
|
||||
@@ -297,9 +374,68 @@ class Client(object):
|
||||
"uuid": uuid,
|
||||
}
|
||||
|
||||
return self._enqueue(msg)
|
||||
return self._enqueue(msg, disable_geoip)
|
||||
|
||||
def _enqueue(self, msg):
|
||||
def capture_exception(
|
||||
self,
|
||||
exception=None,
|
||||
distinct_id=None,
|
||||
properties=None,
|
||||
context=None,
|
||||
timestamp=None,
|
||||
uuid=None,
|
||||
groups=None,
|
||||
):
|
||||
# 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)`"""
|
||||
|
||||
if self.disabled:
|
||||
@@ -307,7 +443,7 @@ 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)
|
||||
@@ -327,6 +463,15 @@ class Client(object):
|
||||
msg["properties"]["$lib"] = "posthog-python"
|
||||
msg["properties"]["$lib_version"] = VERSION
|
||||
|
||||
if disable_geoip is None:
|
||||
disable_geoip = self.disable_geoip
|
||||
|
||||
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)
|
||||
@@ -338,7 +483,14 @@ class Client(object):
|
||||
|
||||
if self.sync_mode:
|
||||
self.log.debug("enqueued with blocking %s.", msg["event"])
|
||||
batch_post(self.api_key, self.host, gzip=self.gzip, timeout=self.timeout, batch=[msg])
|
||||
batch_post(
|
||||
self.api_key,
|
||||
self.host,
|
||||
gzip=self.gzip,
|
||||
timeout=self.timeout,
|
||||
batch=[msg],
|
||||
historical_migration=self.historical_migration,
|
||||
)
|
||||
|
||||
return True, msg
|
||||
|
||||
@@ -378,6 +530,9 @@ class Client(object):
|
||||
self.flush()
|
||||
self.join()
|
||||
|
||||
if self.exception_capture:
|
||||
self.exception_capture.close()
|
||||
|
||||
def _load_feature_flags(self):
|
||||
try:
|
||||
response = get(
|
||||
@@ -415,7 +570,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:
|
||||
@@ -428,7 +583,16 @@ class Client(object):
|
||||
self.poller = Poller(interval=timedelta(seconds=self.poll_interval), execute=self._load_feature_flags)
|
||||
self.poller.start()
|
||||
|
||||
def _compute_flag_locally(self, feature_flag, distinct_id, *, groups={}, person_properties={}, group_properties={}):
|
||||
def _compute_flag_locally(
|
||||
self,
|
||||
feature_flag,
|
||||
distinct_id,
|
||||
*,
|
||||
groups={},
|
||||
person_properties={},
|
||||
group_properties={},
|
||||
warn_on_unknown_groups=True,
|
||||
):
|
||||
if feature_flag.get("ensure_experience_continuity", False):
|
||||
raise InconclusiveMatchError("Flag has experience continuity enabled")
|
||||
|
||||
@@ -450,9 +614,14 @@ class Client(object):
|
||||
if group_name not in groups:
|
||||
# Group flags are never enabled in `groups` aren't passed in
|
||||
# don't failover to `/decide/`, since response will be the same
|
||||
self.log.warning(
|
||||
f"[FEATURE FLAGS] Can't compute group feature flag: {feature_flag['key']} without group names passed in"
|
||||
)
|
||||
if warn_on_unknown_groups:
|
||||
self.log.warning(
|
||||
f"[FEATURE FLAGS] Can't compute group feature flag: {feature_flag['key']} without group names passed in"
|
||||
)
|
||||
else:
|
||||
self.log.debug(
|
||||
f"[FEATURE FLAGS] Can't compute group feature flag: {feature_flag['key']} without group names passed in"
|
||||
)
|
||||
return False
|
||||
|
||||
focused_group_properties = group_properties[group_name]
|
||||
@@ -470,6 +639,7 @@ class Client(object):
|
||||
group_properties={},
|
||||
only_evaluate_locally=False,
|
||||
send_feature_flag_events=True,
|
||||
disable_geoip=None,
|
||||
):
|
||||
response = self.get_feature_flag(
|
||||
key,
|
||||
@@ -479,6 +649,7 @@ class Client(object):
|
||||
group_properties=group_properties,
|
||||
only_evaluate_locally=only_evaluate_locally,
|
||||
send_feature_flag_events=send_feature_flag_events,
|
||||
disable_geoip=disable_geoip,
|
||||
)
|
||||
|
||||
if response is None:
|
||||
@@ -495,11 +666,19 @@ class Client(object):
|
||||
group_properties={},
|
||||
only_evaluate_locally=False,
|
||||
send_feature_flag_events=True,
|
||||
disable_geoip=None,
|
||||
):
|
||||
require("key", key, string_types)
|
||||
require("distinct_id", distinct_id, ID_TYPES)
|
||||
require("groups", groups, dict)
|
||||
|
||||
if self.disabled:
|
||||
return None
|
||||
|
||||
person_properties, group_properties = self._add_local_person_and_group_properties(
|
||||
distinct_id, groups, person_properties, group_properties
|
||||
)
|
||||
|
||||
if self.feature_flags is None and self.personal_api_key:
|
||||
self.load_feature_flags()
|
||||
response = None
|
||||
@@ -528,7 +707,11 @@ class Client(object):
|
||||
if not flag_was_locally_evaluated and not only_evaluate_locally:
|
||||
try:
|
||||
feature_flags = self.get_feature_variants(
|
||||
distinct_id, groups=groups, person_properties=person_properties, group_properties=group_properties
|
||||
distinct_id,
|
||||
groups=groups,
|
||||
person_properties=person_properties,
|
||||
group_properties=group_properties,
|
||||
disable_geoip=disable_geoip,
|
||||
)
|
||||
response = feature_flags.get(key)
|
||||
if response is None:
|
||||
@@ -549,8 +732,10 @@ class Client(object):
|
||||
"$feature_flag": key,
|
||||
"$feature_flag_response": response,
|
||||
"locally_evaluated": flag_was_locally_evaluated,
|
||||
f"$feature/{key}": response,
|
||||
},
|
||||
groups=groups,
|
||||
disable_geoip=disable_geoip,
|
||||
)
|
||||
self.distinct_ids_feature_flags_reported[distinct_id].add(feature_flag_reported_key)
|
||||
return response
|
||||
@@ -566,7 +751,11 @@ class Client(object):
|
||||
group_properties={},
|
||||
only_evaluate_locally=False,
|
||||
send_feature_flag_events=True,
|
||||
disable_geoip=None,
|
||||
):
|
||||
if self.disabled:
|
||||
return None
|
||||
|
||||
if match_value is None:
|
||||
match_value = self.get_feature_flag(
|
||||
key,
|
||||
@@ -574,20 +763,52 @@ 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)
|
||||
response = decide_payloads.get(str(key).lower(), None)
|
||||
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).lower(), None)
|
||||
except Exception as e:
|
||||
self.log.exception(f"[FEATURE FLAGS] Unable to get feature flags and payloads: {e}")
|
||||
|
||||
return response
|
||||
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,
|
||||
)
|
||||
self.distinct_ids_feature_flags_reported[distinct_id].add(feature_flag_reported_key)
|
||||
|
||||
return payload
|
||||
|
||||
def _compute_payload_locally(self, key, match_value):
|
||||
payload = None
|
||||
@@ -602,7 +823,14 @@ class Client(object):
|
||||
return payload
|
||||
|
||||
def get_all_flags(
|
||||
self, distinct_id, *, groups={}, person_properties={}, group_properties={}, only_evaluate_locally=False
|
||||
self,
|
||||
distinct_id,
|
||||
*,
|
||||
groups={},
|
||||
person_properties={},
|
||||
group_properties={},
|
||||
only_evaluate_locally=False,
|
||||
disable_geoip=None,
|
||||
):
|
||||
flags = self.get_all_flags_and_payloads(
|
||||
distinct_id,
|
||||
@@ -610,12 +838,27 @@ class Client(object):
|
||||
person_properties=person_properties,
|
||||
group_properties=group_properties,
|
||||
only_evaluate_locally=only_evaluate_locally,
|
||||
disable_geoip=disable_geoip,
|
||||
)
|
||||
return flags["featureFlags"]
|
||||
|
||||
def get_all_flags_and_payloads(
|
||||
self, distinct_id, *, groups={}, person_properties={}, group_properties={}, only_evaluate_locally=False
|
||||
self,
|
||||
distinct_id,
|
||||
*,
|
||||
groups={},
|
||||
person_properties={},
|
||||
group_properties={},
|
||||
only_evaluate_locally=False,
|
||||
disable_geoip=None,
|
||||
):
|
||||
if self.disabled:
|
||||
return {"featureFlags": None, "featureFlagPayloads": None}
|
||||
|
||||
person_properties, group_properties = self._add_local_person_and_group_properties(
|
||||
distinct_id, groups, person_properties, group_properties
|
||||
)
|
||||
|
||||
flags, payloads, fallback_to_decide = self._get_all_flags_and_payloads_locally(
|
||||
distinct_id, groups=groups, person_properties=person_properties, group_properties=group_properties
|
||||
)
|
||||
@@ -624,7 +867,11 @@ class Client(object):
|
||||
if fallback_to_decide and not only_evaluate_locally:
|
||||
try:
|
||||
flags_and_payloads = self.get_decide(
|
||||
distinct_id, groups=groups, person_properties=person_properties, group_properties=group_properties
|
||||
distinct_id,
|
||||
groups=groups,
|
||||
person_properties=person_properties,
|
||||
group_properties=group_properties,
|
||||
disable_geoip=disable_geoip,
|
||||
)
|
||||
response = flags_and_payloads
|
||||
except Exception as e:
|
||||
@@ -632,7 +879,9 @@ class Client(object):
|
||||
|
||||
return response
|
||||
|
||||
def _get_all_flags_and_payloads_locally(self, distinct_id, *, groups={}, person_properties={}, group_properties={}):
|
||||
def _get_all_flags_and_payloads_locally(
|
||||
self, distinct_id, *, groups={}, person_properties={}, group_properties={}, warn_on_unknown_groups=False
|
||||
):
|
||||
require("distinct_id", distinct_id, ID_TYPES)
|
||||
require("groups", groups, dict)
|
||||
|
||||
@@ -652,6 +901,7 @@ class Client(object):
|
||||
groups=groups,
|
||||
person_properties=person_properties,
|
||||
group_properties=group_properties,
|
||||
warn_on_unknown_groups=warn_on_unknown_groups,
|
||||
)
|
||||
matched_payload = self._compute_payload_locally(flag["key"], flags[flag["key"]])
|
||||
if matched_payload:
|
||||
@@ -667,6 +917,22 @@ class Client(object):
|
||||
|
||||
return flags, payloads, fallback_to_decide
|
||||
|
||||
def feature_flag_definitions(self):
|
||||
return self.feature_flags
|
||||
|
||||
def _add_local_person_and_group_properties(self, distinct_id, groups, person_properties, group_properties):
|
||||
all_person_properties = {"distinct_id": distinct_id, **(person_properties or {})}
|
||||
|
||||
all_group_properties = {}
|
||||
if groups:
|
||||
for group_name in groups:
|
||||
all_group_properties[group_name] = {
|
||||
"$group_key": groups[group_name],
|
||||
**(group_properties.get(group_name) or {}),
|
||||
}
|
||||
|
||||
return all_person_properties, all_group_properties
|
||||
|
||||
|
||||
def require(name, field, data_type):
|
||||
"""Require that the named `field` has the right `data_type`"""
|
||||
|
||||
+16
-6
@@ -12,11 +12,12 @@ try:
|
||||
except ImportError:
|
||||
from Queue import Empty
|
||||
|
||||
MAX_MSG_SIZE = 32 << 10
|
||||
|
||||
# Our servers only accept batches less than 500KB. Here limit is set slightly
|
||||
# lower to leave space for extra data that will be added later, eg. "sentAt".
|
||||
BATCH_SIZE_LIMIT = 475000
|
||||
MAX_MSG_SIZE = 900 * 1024 # 900KiB per event
|
||||
|
||||
# The maximum request body size is currently 20MiB, let's be conservative
|
||||
# in case we want to lower it in the future.
|
||||
BATCH_SIZE_LIMIT = 5 * 1024 * 1024
|
||||
|
||||
|
||||
class Consumer(Thread):
|
||||
@@ -35,6 +36,7 @@ class Consumer(Thread):
|
||||
gzip=False,
|
||||
retries=10,
|
||||
timeout=15,
|
||||
historical_migration=False,
|
||||
):
|
||||
"""Create a consumer thread."""
|
||||
Thread.__init__(self)
|
||||
@@ -54,6 +56,7 @@ class Consumer(Thread):
|
||||
self.running = True
|
||||
self.retries = retries
|
||||
self.timeout = timeout
|
||||
self.historical_migration = historical_migration
|
||||
|
||||
def run(self):
|
||||
"""Runs the consumer."""
|
||||
@@ -104,7 +107,7 @@ class Consumer(Thread):
|
||||
item = queue.get(block=True, timeout=self.flush_interval - elapsed)
|
||||
item_size = len(json.dumps(item, cls=DatetimeSerializer).encode())
|
||||
if item_size > MAX_MSG_SIZE:
|
||||
self.log.error("Item exceeds 32kb limit, dropping. (%s)", str(item))
|
||||
self.log.error("Item exceeds 900kib limit, dropping. (%s)", str(item))
|
||||
continue
|
||||
items.append(item)
|
||||
total_size += item_size
|
||||
@@ -133,6 +136,13 @@ class Consumer(Thread):
|
||||
|
||||
@backoff.on_exception(backoff.expo, Exception, max_tries=self.retries + 1, giveup=fatal_exception)
|
||||
def send_request():
|
||||
batch_post(self.api_key, self.host, gzip=self.gzip, timeout=self.timeout, batch=batch)
|
||||
batch_post(
|
||||
self.api_key,
|
||||
self.host,
|
||||
gzip=self.gzip,
|
||||
timeout=self.timeout,
|
||||
batch=batch,
|
||||
historical_migration=self.historical_migration,
|
||||
)
|
||||
|
||||
send_request()
|
||||
|
||||
@@ -0,0 +1,64 @@
|
||||
import logging
|
||||
import sys
|
||||
import threading
|
||||
from enum import Enum
|
||||
from typing import TYPE_CHECKING, List, Optional
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from posthog.client import Client
|
||||
|
||||
|
||||
class Integrations(str, Enum):
|
||||
Django = "django"
|
||||
|
||||
|
||||
class ExceptionCapture:
|
||||
# TODO: Add client side rate limiting to prevent spamming the server with exceptions
|
||||
|
||||
log = logging.getLogger("posthog")
|
||||
|
||||
def __init__(self, client: "Client", integrations: Optional[List[Integrations]] = None):
|
||||
self.client = client
|
||||
self.original_excepthook = sys.excepthook
|
||||
sys.excepthook = self.exception_handler
|
||||
threading.excepthook = self.thread_exception_handler
|
||||
self.enabled_integrations = []
|
||||
|
||||
for integration in integrations or []:
|
||||
# TODO: Maybe find a better way of enabling integrations
|
||||
# This is very annoying currently if we had to add any configuration per integration
|
||||
if integration == Integrations.Django:
|
||||
try:
|
||||
from posthog.exception_integrations.django import DjangoIntegration
|
||||
|
||||
enabled_integration = DjangoIntegration(self.exception_receiver)
|
||||
self.enabled_integrations.append(enabled_integration)
|
||||
except Exception as e:
|
||||
self.log.exception(f"Failed to enable Django integration: {e}")
|
||||
|
||||
def close(self):
|
||||
sys.excepthook = self.original_excepthook
|
||||
for integration in self.enabled_integrations:
|
||||
integration.uninstall()
|
||||
|
||||
def exception_handler(self, exc_type, exc_value, exc_traceback):
|
||||
# don't affect default behaviour.
|
||||
self.capture_exception((exc_type, exc_value, exc_traceback))
|
||||
self.original_excepthook(exc_type, exc_value, exc_traceback)
|
||||
|
||||
def thread_exception_handler(self, args):
|
||||
self.capture_exception((args.exc_type, args.exc_value, args.exc_traceback))
|
||||
|
||||
def exception_receiver(self, exc_info, extra_properties):
|
||||
if "distinct_id" in extra_properties:
|
||||
metadata = {"distinct_id": extra_properties["distinct_id"]}
|
||||
else:
|
||||
metadata = None
|
||||
self.capture_exception((exc_info[0], exc_info[1], exc_info[2]), metadata)
|
||||
|
||||
def capture_exception(self, exception, metadata=None):
|
||||
try:
|
||||
distinct_id = metadata.get("distinct_id") if metadata else None
|
||||
self.client.capture_exception(exception, distinct_id)
|
||||
except Exception as e:
|
||||
self.log.exception(f"Failed to capture exception: {e}")
|
||||
@@ -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`.
|
||||
"""
|
||||
@@ -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)
|
||||
@@ -0,0 +1,873 @@
|
||||
# copied and adapted from https://github.com/getsentry/sentry-python/blob/269d96d6e9821122fbff280e6a26956e5ed03c0b/sentry_sdk/utils.py#L689
|
||||
# 💖open source (under MIT License)
|
||||
# We want to keep payloads as similar to Sentry as possible for easy interoperability
|
||||
|
||||
import linecache
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
try:
|
||||
# Python 3.11
|
||||
from builtins import BaseExceptionGroup
|
||||
except ImportError:
|
||||
# Python 3.10 and below
|
||||
BaseExceptionGroup = None # type: ignore
|
||||
|
||||
|
||||
DEFAULT_MAX_VALUE_LENGTH = 1024
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
||||
from types import FrameType, TracebackType
|
||||
from typing import ( # noqa: F401
|
||||
Any,
|
||||
Callable,
|
||||
Dict,
|
||||
Iterator,
|
||||
List,
|
||||
Literal,
|
||||
Optional,
|
||||
Set,
|
||||
Tuple,
|
||||
Type,
|
||||
TypedDict,
|
||||
TypeVar,
|
||||
Union,
|
||||
cast,
|
||||
)
|
||||
|
||||
ExcInfo = Union[
|
||||
Tuple[Type[BaseException], BaseException, Optional[TracebackType]],
|
||||
Tuple[None, None, None],
|
||||
]
|
||||
LogLevelStr = Literal["fatal", "critical", "error", "warning", "info", "debug"]
|
||||
|
||||
Event = TypedDict(
|
||||
"Event",
|
||||
{
|
||||
"breadcrumbs": Dict[Literal["values"], List[Dict[str, Any]]], # TODO: We can expand on this type
|
||||
"check_in_id": str,
|
||||
"contexts": Dict[str, Dict[str, object]],
|
||||
"dist": str,
|
||||
"duration": Optional[float],
|
||||
"environment": str,
|
||||
"errors": List[Dict[str, Any]], # TODO: We can expand on this type
|
||||
"event_id": str,
|
||||
"exception": Dict[Literal["values"], List[Dict[str, Any]]], # TODO: We can expand on this type
|
||||
# "extra": MutableMapping[str, object],
|
||||
# "fingerprint": List[str],
|
||||
"level": LogLevelStr,
|
||||
# "logentry": Mapping[str, object],
|
||||
"logger": str,
|
||||
# "measurements": Dict[str, MeasurementValue],
|
||||
"message": str,
|
||||
"modules": Dict[str, str],
|
||||
# "monitor_config": Mapping[str, object],
|
||||
"monitor_slug": Optional[str],
|
||||
"platform": Literal["python"],
|
||||
"profile": object, # Should be sentry_sdk.profiler.Profile, but we can't import that here due to circular imports
|
||||
"release": str,
|
||||
"request": Dict[str, object],
|
||||
# "sdk": Mapping[str, object],
|
||||
"server_name": str,
|
||||
"spans": List[Dict[str, object]],
|
||||
"stacktrace": Dict[str, object], # We access this key in the code, but I am unsure whether we ever set it
|
||||
"start_timestamp": datetime,
|
||||
"status": Optional[str],
|
||||
# "tags": MutableMapping[
|
||||
# str, str
|
||||
# ], # Tags must be less than 200 characters each
|
||||
"threads": Dict[Literal["values"], List[Dict[str, Any]]], # TODO: We can expand on this type
|
||||
"timestamp": Optional[datetime], # Must be set before sending the event
|
||||
"transaction": str,
|
||||
# "transaction_info": Mapping[str, Any], # TODO: We can expand on this type
|
||||
"type": Literal["check_in", "transaction"],
|
||||
"user": Dict[str, object],
|
||||
"_metrics_summary": Dict[str, object],
|
||||
},
|
||||
total=False,
|
||||
)
|
||||
|
||||
|
||||
epoch = datetime(1970, 1, 1)
|
||||
|
||||
|
||||
BASE64_ALPHABET = re.compile(r"^[a-zA-Z0-9/+=]*$")
|
||||
|
||||
SENSITIVE_DATA_SUBSTITUTE = "[Filtered]"
|
||||
|
||||
|
||||
def to_timestamp(value):
|
||||
# type: (datetime) -> float
|
||||
return (value - epoch).total_seconds()
|
||||
|
||||
|
||||
def format_timestamp(value):
|
||||
# type: (datetime) -> str
|
||||
return value.strftime("%Y-%m-%dT%H:%M:%S.%fZ")
|
||||
|
||||
|
||||
def event_hint_with_exc_info(exc_info=None):
|
||||
# type: (Optional[ExcInfo]) -> Dict[str, Optional[ExcInfo]]
|
||||
"""Creates a hint with the exc info filled in."""
|
||||
if exc_info is None:
|
||||
exc_info = sys.exc_info()
|
||||
else:
|
||||
exc_info = exc_info_from_error(exc_info)
|
||||
if exc_info[0] is None:
|
||||
exc_info = None
|
||||
return {"exc_info": exc_info}
|
||||
|
||||
|
||||
class AnnotatedValue:
|
||||
"""
|
||||
Meta information for a data field in the event payload.
|
||||
This is to tell Relay that we have tampered with the fields value.
|
||||
See:
|
||||
https://github.com/getsentry/relay/blob/be12cd49a0f06ea932ed9b9f93a655de5d6ad6d1/relay-general/src/types/meta.rs#L407-L423
|
||||
"""
|
||||
|
||||
__slots__ = ("value", "metadata")
|
||||
|
||||
def __init__(self, value, metadata):
|
||||
# type: (Optional[Any], Dict[str, Any]) -> None
|
||||
self.value = value
|
||||
self.metadata = metadata
|
||||
|
||||
def __eq__(self, other):
|
||||
# type: (Any) -> bool
|
||||
if not isinstance(other, AnnotatedValue):
|
||||
return False
|
||||
|
||||
return self.value == other.value and self.metadata == other.metadata
|
||||
|
||||
@classmethod
|
||||
def removed_because_raw_data(cls):
|
||||
# type: () -> AnnotatedValue
|
||||
"""The value was removed because it could not be parsed. This is done for request body values that are not json nor a form."""
|
||||
return AnnotatedValue(
|
||||
value="",
|
||||
metadata={
|
||||
"rem": [ # Remark
|
||||
[
|
||||
"!raw", # Unparsable raw data
|
||||
"x", # The fields original value was removed
|
||||
]
|
||||
]
|
||||
},
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def removed_because_over_size_limit(cls):
|
||||
# type: () -> AnnotatedValue
|
||||
"""The actual value was removed because the size of the field exceeded the configured maximum size (specified with the max_request_body_size sdk option)"""
|
||||
return AnnotatedValue(
|
||||
value="",
|
||||
metadata={
|
||||
"rem": [ # Remark
|
||||
[
|
||||
"!config", # Because of configured maximum size
|
||||
"x", # The fields original value was removed
|
||||
]
|
||||
]
|
||||
},
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def substituted_because_contains_sensitive_data(cls):
|
||||
# type: () -> AnnotatedValue
|
||||
"""The actual value was removed because it contained sensitive information."""
|
||||
return AnnotatedValue(
|
||||
value=SENSITIVE_DATA_SUBSTITUTE,
|
||||
metadata={
|
||||
"rem": [ # Remark
|
||||
[
|
||||
"!config", # Because of SDK configuration (in this case the config is the hard coded removal of certain django cookies)
|
||||
"s", # The fields original value was substituted
|
||||
]
|
||||
]
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
T = TypeVar("T")
|
||||
Annotated = Union[AnnotatedValue, T]
|
||||
|
||||
|
||||
def get_type_name(cls):
|
||||
# type: (Optional[type]) -> Optional[str]
|
||||
return getattr(cls, "__qualname__", None) or getattr(cls, "__name__", None)
|
||||
|
||||
|
||||
def get_type_module(cls):
|
||||
# type: (Optional[type]) -> Optional[str]
|
||||
mod = getattr(cls, "__module__", None)
|
||||
if mod not in (None, "builtins", "__builtins__"):
|
||||
return mod
|
||||
return None
|
||||
|
||||
|
||||
def should_hide_frame(frame: "FrameType") -> bool:
|
||||
try:
|
||||
mod = frame.f_globals["__name__"]
|
||||
if mod.startswith("sentry_sdk."):
|
||||
return True
|
||||
except (AttributeError, KeyError):
|
||||
pass
|
||||
|
||||
for flag_name in "__traceback_hide__", "__tracebackhide__":
|
||||
try:
|
||||
if frame.f_locals[flag_name]:
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def iter_stacks(tb):
|
||||
# type: (Optional[TracebackType]) -> Iterator[TracebackType]
|
||||
tb_ = tb # type: Optional[TracebackType]
|
||||
while tb_ is not None:
|
||||
if not should_hide_frame(tb_.tb_frame):
|
||||
yield tb_
|
||||
tb_ = tb_.tb_next
|
||||
|
||||
|
||||
def get_lines_from_file(
|
||||
filename, # type: str
|
||||
lineno, # type: int
|
||||
max_length=None, # type: Optional[int]
|
||||
loader=None, # type: Optional[Any]
|
||||
module=None, # type: Optional[str]
|
||||
):
|
||||
# type: (...) -> Tuple[List[Annotated[str]], Optional[Annotated[str]], List[Annotated[str]]]
|
||||
context_lines = 5
|
||||
source = None
|
||||
if loader is not None and hasattr(loader, "get_source"):
|
||||
try:
|
||||
source_str = loader.get_source(module) # type: Optional[str]
|
||||
except (ImportError, IOError):
|
||||
source_str = None
|
||||
if source_str is not None:
|
||||
source = source_str.splitlines()
|
||||
|
||||
if source is None:
|
||||
try:
|
||||
source = linecache.getlines(filename)
|
||||
except (OSError, IOError):
|
||||
return [], None, []
|
||||
|
||||
if not source:
|
||||
return [], None, []
|
||||
|
||||
lower_bound = max(0, lineno - context_lines)
|
||||
upper_bound = min(lineno + 1 + context_lines, len(source))
|
||||
|
||||
try:
|
||||
pre_context = [strip_string(line.strip("\r\n"), max_length=max_length) for line in source[lower_bound:lineno]]
|
||||
context_line = strip_string(source[lineno].strip("\r\n"), max_length=max_length)
|
||||
post_context = [
|
||||
strip_string(line.strip("\r\n"), max_length=max_length)
|
||||
for line in source[(lineno + 1) : upper_bound] # noqa: E203
|
||||
]
|
||||
return pre_context, context_line, post_context
|
||||
except IndexError:
|
||||
# the file may have changed since it was loaded into memory
|
||||
return [], None, []
|
||||
|
||||
|
||||
def get_source_context(
|
||||
frame, # type: FrameType
|
||||
tb_lineno, # type: int
|
||||
max_value_length=None, # type: Optional[int]
|
||||
):
|
||||
# type: (...) -> Tuple[List[Annotated[str]], Optional[Annotated[str]], List[Annotated[str]]]
|
||||
try:
|
||||
abs_path = frame.f_code.co_filename # type: Optional[str]
|
||||
except Exception:
|
||||
abs_path = None
|
||||
try:
|
||||
module = frame.f_globals["__name__"]
|
||||
except Exception:
|
||||
return [], None, []
|
||||
try:
|
||||
loader = frame.f_globals["__loader__"]
|
||||
except Exception:
|
||||
loader = None
|
||||
lineno = tb_lineno - 1
|
||||
if lineno is not None and abs_path:
|
||||
return get_lines_from_file(abs_path, lineno, max_value_length, loader=loader, module=module)
|
||||
return [], None, []
|
||||
|
||||
|
||||
def safe_str(value):
|
||||
# type: (Any) -> str
|
||||
try:
|
||||
return str(value)
|
||||
except Exception:
|
||||
return safe_repr(value)
|
||||
|
||||
|
||||
def safe_repr(value):
|
||||
# type: (Any) -> str
|
||||
try:
|
||||
return repr(value)
|
||||
except Exception:
|
||||
return "<broken repr>"
|
||||
|
||||
|
||||
def filename_for_module(module, abs_path):
|
||||
# type: (Optional[str], Optional[str]) -> Optional[str]
|
||||
if not abs_path or not module:
|
||||
return abs_path
|
||||
|
||||
try:
|
||||
if abs_path.endswith(".pyc"):
|
||||
abs_path = abs_path[:-1]
|
||||
|
||||
base_module = module.split(".", 1)[0]
|
||||
if base_module == module:
|
||||
return os.path.basename(abs_path)
|
||||
|
||||
base_module_path = sys.modules[base_module].__file__
|
||||
if not base_module_path:
|
||||
return abs_path
|
||||
|
||||
return abs_path.split(base_module_path.rsplit(os.sep, 2)[0], 1)[-1].lstrip(os.sep)
|
||||
except Exception:
|
||||
return abs_path
|
||||
|
||||
|
||||
def serialize_frame(
|
||||
frame,
|
||||
tb_lineno=None,
|
||||
include_local_variables=True,
|
||||
include_source_context=True,
|
||||
max_value_length=None,
|
||||
custom_repr=None,
|
||||
):
|
||||
# type: (FrameType, Optional[int], bool, bool, Optional[int], Optional[Callable[..., Optional[str]]]) -> Dict[str, Any]
|
||||
f_code = getattr(frame, "f_code", None)
|
||||
if not f_code:
|
||||
abs_path = None
|
||||
function = None
|
||||
else:
|
||||
abs_path = frame.f_code.co_filename
|
||||
function = frame.f_code.co_name
|
||||
try:
|
||||
module = frame.f_globals["__name__"]
|
||||
except Exception:
|
||||
module = None
|
||||
|
||||
if tb_lineno is None:
|
||||
tb_lineno = frame.f_lineno
|
||||
|
||||
rv = {
|
||||
"platform": "python",
|
||||
"filename": filename_for_module(module, abs_path) or None,
|
||||
"abs_path": os.path.abspath(abs_path) if abs_path else None,
|
||||
"function": function or "<unknown>",
|
||||
"module": module,
|
||||
"lineno": tb_lineno,
|
||||
} # type: Dict[str, Any]
|
||||
|
||||
if include_source_context:
|
||||
rv["pre_context"], rv["context_line"], rv["post_context"] = get_source_context(
|
||||
frame, tb_lineno, max_value_length
|
||||
)
|
||||
|
||||
if include_local_variables:
|
||||
# TODO(nk): Sort out this current invalid import
|
||||
# from sentry_sdk.serializer import serialize
|
||||
|
||||
# rv["vars"] = serialize(
|
||||
# dict(frame.f_locals), is_vars=True, custom_repr=custom_repr
|
||||
# )
|
||||
pass
|
||||
|
||||
return rv
|
||||
|
||||
|
||||
def current_stacktrace(
|
||||
include_local_variables=True, # type: bool
|
||||
include_source_context=True, # type: bool
|
||||
max_value_length=None, # type: Optional[int]
|
||||
):
|
||||
# type: (...) -> Dict[str, Any]
|
||||
__tracebackhide__ = True
|
||||
frames = []
|
||||
|
||||
f = sys._getframe() # type: Optional[FrameType]
|
||||
while f is not None:
|
||||
if not should_hide_frame(f):
|
||||
frames.append(
|
||||
serialize_frame(
|
||||
f,
|
||||
include_local_variables=include_local_variables,
|
||||
include_source_context=include_source_context,
|
||||
max_value_length=max_value_length,
|
||||
)
|
||||
)
|
||||
f = f.f_back
|
||||
|
||||
frames.reverse()
|
||||
|
||||
return {"frames": frames, "type": "raw"}
|
||||
|
||||
|
||||
def get_errno(exc_value):
|
||||
# type: (BaseException) -> Optional[Any]
|
||||
return getattr(exc_value, "errno", None)
|
||||
|
||||
|
||||
def get_error_message(exc_value):
|
||||
# type: (Optional[BaseException]) -> str
|
||||
return getattr(exc_value, "message", "") or getattr(exc_value, "detail", "") or safe_str(exc_value)
|
||||
|
||||
|
||||
def single_exception_from_error_tuple(
|
||||
exc_type, # type: Optional[type]
|
||||
exc_value, # type: Optional[BaseException]
|
||||
tb, # type: Optional[TracebackType]
|
||||
client_options=None, # type: Optional[Dict[str, Any]]
|
||||
mechanism=None, # type: Optional[Dict[str, Any]]
|
||||
exception_id=None, # type: Optional[int]
|
||||
parent_id=None, # type: Optional[int]
|
||||
source=None, # type: Optional[str]
|
||||
):
|
||||
# type: (...) -> Dict[str, Any]
|
||||
"""
|
||||
Creates a dict that goes into the events `exception.values` list and is ingestible by Sentry.
|
||||
|
||||
See the Exception Interface documentation for more details:
|
||||
https://develop.sentry.dev/sdk/event-payloads/exception/
|
||||
"""
|
||||
exception_value = {} # type: Dict[str, Any]
|
||||
exception_value["mechanism"] = mechanism.copy() if mechanism else {"type": "generic", "handled": True}
|
||||
if exception_id is not None:
|
||||
exception_value["mechanism"]["exception_id"] = exception_id
|
||||
|
||||
if exc_value is not None:
|
||||
errno = get_errno(exc_value)
|
||||
else:
|
||||
errno = None
|
||||
|
||||
if errno is not None:
|
||||
exception_value["mechanism"].setdefault("meta", {}).setdefault("errno", {}).setdefault("number", errno)
|
||||
|
||||
if source is not None:
|
||||
exception_value["mechanism"]["source"] = source
|
||||
|
||||
is_root_exception = exception_id == 0
|
||||
if not is_root_exception and parent_id is not None:
|
||||
exception_value["mechanism"]["parent_id"] = parent_id
|
||||
exception_value["mechanism"]["type"] = "chained"
|
||||
|
||||
if is_root_exception and "type" not in exception_value["mechanism"]:
|
||||
exception_value["mechanism"]["type"] = "generic"
|
||||
|
||||
is_exception_group = BaseExceptionGroup is not None and isinstance(exc_value, BaseExceptionGroup)
|
||||
if is_exception_group:
|
||||
exception_value["mechanism"]["is_exception_group"] = True
|
||||
|
||||
exception_value["module"] = get_type_module(exc_type)
|
||||
exception_value["type"] = get_type_name(exc_type)
|
||||
exception_value["value"] = get_error_message(exc_value)
|
||||
|
||||
if client_options is None:
|
||||
include_local_variables = True
|
||||
include_source_context = True
|
||||
max_value_length = DEFAULT_MAX_VALUE_LENGTH # fallback
|
||||
custom_repr = None
|
||||
else:
|
||||
include_local_variables = client_options["include_local_variables"]
|
||||
include_source_context = client_options["include_source_context"]
|
||||
max_value_length = client_options["max_value_length"]
|
||||
custom_repr = client_options.get("custom_repr")
|
||||
|
||||
frames = [
|
||||
serialize_frame(
|
||||
tb.tb_frame,
|
||||
tb_lineno=tb.tb_lineno,
|
||||
include_local_variables=include_local_variables,
|
||||
include_source_context=include_source_context,
|
||||
max_value_length=max_value_length,
|
||||
custom_repr=custom_repr,
|
||||
)
|
||||
for tb in iter_stacks(tb)
|
||||
]
|
||||
|
||||
if frames:
|
||||
exception_value["stacktrace"] = {"frames": frames, "type": "raw"}
|
||||
|
||||
return exception_value
|
||||
|
||||
|
||||
HAS_CHAINED_EXCEPTIONS = hasattr(Exception, "__suppress_context__")
|
||||
|
||||
if HAS_CHAINED_EXCEPTIONS:
|
||||
|
||||
def walk_exception_chain(exc_info):
|
||||
# type: (ExcInfo) -> Iterator[ExcInfo]
|
||||
exc_type, exc_value, tb = exc_info
|
||||
|
||||
seen_exceptions = []
|
||||
seen_exception_ids = set() # type: Set[int]
|
||||
|
||||
while exc_type is not None and exc_value is not None and id(exc_value) not in seen_exception_ids:
|
||||
yield exc_type, exc_value, tb
|
||||
|
||||
# Avoid hashing random types we don't know anything
|
||||
# about. Use the list to keep a ref so that the `id` is
|
||||
# not used for another object.
|
||||
seen_exceptions.append(exc_value)
|
||||
seen_exception_ids.add(id(exc_value))
|
||||
|
||||
if exc_value.__suppress_context__:
|
||||
cause = exc_value.__cause__
|
||||
else:
|
||||
cause = exc_value.__context__
|
||||
if cause is None:
|
||||
break
|
||||
exc_type = type(cause)
|
||||
exc_value = cause
|
||||
tb = getattr(cause, "__traceback__", None)
|
||||
|
||||
else:
|
||||
|
||||
def walk_exception_chain(exc_info):
|
||||
# type: (ExcInfo) -> Iterator[ExcInfo]
|
||||
yield exc_info
|
||||
|
||||
|
||||
def exceptions_from_error(
|
||||
exc_type, # type: Optional[type]
|
||||
exc_value, # type: Optional[BaseException]
|
||||
tb, # type: Optional[TracebackType]
|
||||
client_options=None, # type: Optional[Dict[str, Any]]
|
||||
mechanism=None, # type: Optional[Dict[str, Any]]
|
||||
exception_id=0, # type: int
|
||||
parent_id=0, # type: int
|
||||
source=None, # type: Optional[str]
|
||||
):
|
||||
# type: (...) -> Tuple[int, List[Dict[str, Any]]]
|
||||
"""
|
||||
Creates the list of exceptions.
|
||||
This can include chained exceptions and exceptions from an ExceptionGroup.
|
||||
|
||||
See the Exception Interface documentation for more details:
|
||||
https://develop.sentry.dev/sdk/event-payloads/exception/
|
||||
"""
|
||||
|
||||
parent = single_exception_from_error_tuple(
|
||||
exc_type=exc_type,
|
||||
exc_value=exc_value,
|
||||
tb=tb,
|
||||
client_options=client_options,
|
||||
mechanism=mechanism,
|
||||
exception_id=exception_id,
|
||||
parent_id=parent_id,
|
||||
source=source,
|
||||
)
|
||||
exceptions = [parent]
|
||||
|
||||
parent_id = exception_id
|
||||
exception_id += 1
|
||||
|
||||
should_supress_context = hasattr(exc_value, "__suppress_context__") and exc_value.__suppress_context__ # type: ignore
|
||||
if should_supress_context:
|
||||
# Add direct cause.
|
||||
# The field `__cause__` is set when raised with the exception (using the `from` keyword).
|
||||
exception_has_cause = exc_value and hasattr(exc_value, "__cause__") and exc_value.__cause__ is not None
|
||||
if exception_has_cause:
|
||||
cause = exc_value.__cause__ # type: ignore
|
||||
(exception_id, child_exceptions) = exceptions_from_error(
|
||||
exc_type=type(cause),
|
||||
exc_value=cause,
|
||||
tb=getattr(cause, "__traceback__", None),
|
||||
client_options=client_options,
|
||||
mechanism=mechanism,
|
||||
exception_id=exception_id,
|
||||
source="__cause__",
|
||||
)
|
||||
exceptions.extend(child_exceptions)
|
||||
|
||||
else:
|
||||
# Add indirect cause.
|
||||
# The field `__context__` is assigned if another exception occurs while handling the exception.
|
||||
exception_has_content = exc_value and hasattr(exc_value, "__context__") and exc_value.__context__ is not None
|
||||
if exception_has_content:
|
||||
context = exc_value.__context__ # type: ignore
|
||||
(exception_id, child_exceptions) = exceptions_from_error(
|
||||
exc_type=type(context),
|
||||
exc_value=context,
|
||||
tb=getattr(context, "__traceback__", None),
|
||||
client_options=client_options,
|
||||
mechanism=mechanism,
|
||||
exception_id=exception_id,
|
||||
source="__context__",
|
||||
)
|
||||
exceptions.extend(child_exceptions)
|
||||
|
||||
# Add exceptions from an ExceptionGroup.
|
||||
is_exception_group = exc_value and hasattr(exc_value, "exceptions")
|
||||
if is_exception_group:
|
||||
for idx, e in enumerate(exc_value.exceptions): # type: ignore
|
||||
(exception_id, child_exceptions) = exceptions_from_error(
|
||||
exc_type=type(e),
|
||||
exc_value=e,
|
||||
tb=getattr(e, "__traceback__", None),
|
||||
client_options=client_options,
|
||||
mechanism=mechanism,
|
||||
exception_id=exception_id,
|
||||
parent_id=parent_id,
|
||||
source="exceptions[%s]" % idx,
|
||||
)
|
||||
exceptions.extend(child_exceptions)
|
||||
|
||||
return (exception_id, exceptions)
|
||||
|
||||
|
||||
def exceptions_from_error_tuple(
|
||||
exc_info, # type: ExcInfo
|
||||
client_options=None, # type: Optional[Dict[str, Any]]
|
||||
mechanism=None, # type: Optional[Dict[str, Any]]
|
||||
):
|
||||
# type: (...) -> List[Dict[str, Any]]
|
||||
exc_type, exc_value, tb = exc_info
|
||||
|
||||
is_exception_group = BaseExceptionGroup is not None and isinstance(exc_value, BaseExceptionGroup)
|
||||
|
||||
if is_exception_group:
|
||||
(_, exceptions) = exceptions_from_error(
|
||||
exc_type=exc_type,
|
||||
exc_value=exc_value,
|
||||
tb=tb,
|
||||
client_options=client_options,
|
||||
mechanism=mechanism,
|
||||
exception_id=0,
|
||||
parent_id=0,
|
||||
)
|
||||
|
||||
else:
|
||||
exceptions = []
|
||||
for exc_type, exc_value, tb in walk_exception_chain(exc_info):
|
||||
exceptions.append(single_exception_from_error_tuple(exc_type, exc_value, tb, client_options, mechanism))
|
||||
|
||||
exceptions.reverse()
|
||||
|
||||
return exceptions
|
||||
|
||||
|
||||
def to_string(value):
|
||||
# type: (str) -> str
|
||||
try:
|
||||
return str(value)
|
||||
except UnicodeDecodeError:
|
||||
return repr(value)[1:-1]
|
||||
|
||||
|
||||
def iter_event_stacktraces(event):
|
||||
# type: (Event) -> Iterator[Dict[str, Any]]
|
||||
if "stacktrace" in event:
|
||||
yield event["stacktrace"]
|
||||
if "threads" in event:
|
||||
for thread in event["threads"].get("values") or ():
|
||||
if "stacktrace" in thread:
|
||||
yield thread["stacktrace"]
|
||||
if "exception" in event:
|
||||
for exception in event["exception"].get("values") or ():
|
||||
if "stacktrace" in exception:
|
||||
yield exception["stacktrace"]
|
||||
|
||||
|
||||
def iter_event_frames(event):
|
||||
# type: (Event) -> Iterator[Dict[str, Any]]
|
||||
for stacktrace in iter_event_stacktraces(event):
|
||||
for frame in stacktrace.get("frames") or ():
|
||||
yield frame
|
||||
|
||||
|
||||
def handle_in_app(event, in_app_exclude=None, in_app_include=None, project_root=None):
|
||||
# type: (Event, Optional[List[str]], Optional[List[str]], Optional[str]) -> Event
|
||||
for stacktrace in iter_event_stacktraces(event):
|
||||
set_in_app_in_frames(
|
||||
stacktrace.get("frames"),
|
||||
in_app_exclude=in_app_exclude,
|
||||
in_app_include=in_app_include,
|
||||
project_root=project_root,
|
||||
)
|
||||
|
||||
return event
|
||||
|
||||
|
||||
def set_in_app_in_frames(frames, in_app_exclude, in_app_include, project_root=None):
|
||||
# type: (Any, Optional[List[str]], Optional[List[str]], Optional[str]) -> Optional[Any]
|
||||
if not frames:
|
||||
return None
|
||||
|
||||
for frame in frames:
|
||||
# if frame has already been marked as in_app, skip it
|
||||
current_in_app = frame.get("in_app")
|
||||
if current_in_app is not None:
|
||||
continue
|
||||
|
||||
module = frame.get("module")
|
||||
|
||||
# check if module in frame is in the list of modules to include
|
||||
if _module_in_list(module, in_app_include):
|
||||
frame["in_app"] = True
|
||||
continue
|
||||
|
||||
# check if module in frame is in the list of modules to exclude
|
||||
if _module_in_list(module, in_app_exclude):
|
||||
frame["in_app"] = False
|
||||
continue
|
||||
|
||||
# if frame has no abs_path, skip further checks
|
||||
abs_path = frame.get("abs_path")
|
||||
if abs_path is None:
|
||||
continue
|
||||
|
||||
if _is_external_source(abs_path):
|
||||
frame["in_app"] = False
|
||||
continue
|
||||
|
||||
if _is_in_project_root(abs_path, project_root):
|
||||
frame["in_app"] = True
|
||||
continue
|
||||
|
||||
return frames
|
||||
|
||||
|
||||
def exc_info_from_error(error):
|
||||
# type: (Union[BaseException, ExcInfo]) -> ExcInfo
|
||||
if isinstance(error, tuple) and len(error) == 3:
|
||||
exc_type, exc_value, tb = error
|
||||
elif isinstance(error, BaseException):
|
||||
tb = getattr(error, "__traceback__", None)
|
||||
if tb is not None:
|
||||
exc_type = type(error)
|
||||
exc_value = error
|
||||
else:
|
||||
exc_type, exc_value, tb = sys.exc_info()
|
||||
if exc_value is not error:
|
||||
tb = None
|
||||
exc_value = error
|
||||
exc_type = type(error)
|
||||
|
||||
else:
|
||||
raise ValueError("Expected Exception object to report, got %s!" % type(error))
|
||||
|
||||
exc_info = (exc_type, exc_value, tb)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
# This cast is safe because exc_type and exc_value are either both
|
||||
# None or both not None.
|
||||
exc_info = cast(ExcInfo, exc_info)
|
||||
|
||||
return exc_info
|
||||
|
||||
|
||||
def event_from_exception(
|
||||
exc_info, # type: Union[BaseException, ExcInfo]
|
||||
client_options=None, # type: Optional[Dict[str, Any]]
|
||||
mechanism=None, # type: Optional[Dict[str, Any]]
|
||||
):
|
||||
# type: (...) -> Tuple[Event, Dict[str, Any]]
|
||||
exc_info = exc_info_from_error(exc_info)
|
||||
hint = event_hint_with_exc_info(exc_info)
|
||||
return (
|
||||
{
|
||||
"level": "error",
|
||||
"exception": {"values": exceptions_from_error_tuple(exc_info, client_options, mechanism)},
|
||||
},
|
||||
hint,
|
||||
)
|
||||
|
||||
|
||||
def _module_in_list(name, items):
|
||||
# type: (str, Optional[List[str]]) -> bool
|
||||
if name is None:
|
||||
return False
|
||||
|
||||
if not items:
|
||||
return False
|
||||
|
||||
for item in items:
|
||||
if item == name or name.startswith(item + "."):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def _is_external_source(abs_path):
|
||||
# type: (str) -> bool
|
||||
# check if frame is in 'site-packages' or 'dist-packages'
|
||||
external_source = re.search(r"[\\/](?:dist|site)-packages[\\/]", abs_path) is not None
|
||||
return external_source
|
||||
|
||||
|
||||
def _is_in_project_root(abs_path, project_root):
|
||||
# type: (str, Optional[str]) -> bool
|
||||
if project_root is None:
|
||||
return False
|
||||
|
||||
# check if path is in the project root
|
||||
if abs_path.startswith(project_root):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def _truncate_by_bytes(string, max_bytes):
|
||||
# type: (str, int) -> str
|
||||
"""
|
||||
Truncate a UTF-8-encodable string to the last full codepoint so that it fits in max_bytes.
|
||||
"""
|
||||
truncated = string.encode("utf-8")[: max_bytes - 3].decode("utf-8", errors="ignore")
|
||||
|
||||
return truncated + "..."
|
||||
|
||||
|
||||
def _get_size_in_bytes(value):
|
||||
# type: (str) -> Optional[int]
|
||||
try:
|
||||
return len(value.encode("utf-8"))
|
||||
except (UnicodeEncodeError, UnicodeDecodeError):
|
||||
return None
|
||||
|
||||
|
||||
def strip_string(value, max_length=None):
|
||||
# type: (str, Optional[int]) -> Union[AnnotatedValue, str]
|
||||
if not value:
|
||||
return value
|
||||
|
||||
if max_length is None:
|
||||
max_length = DEFAULT_MAX_VALUE_LENGTH
|
||||
|
||||
byte_size = _get_size_in_bytes(value)
|
||||
text_size = len(value)
|
||||
|
||||
if byte_size is not None and byte_size > max_length:
|
||||
# truncate to max_length bytes, preserving code points
|
||||
truncated_value = _truncate_by_bytes(value, max_length)
|
||||
elif text_size is not None and text_size > max_length:
|
||||
# fallback to truncating by string length
|
||||
truncated_value = value[: max_length - 3] + "..."
|
||||
else:
|
||||
return value
|
||||
|
||||
return AnnotatedValue(
|
||||
value=truncated_value,
|
||||
metadata={
|
||||
"len": byte_size or text_size,
|
||||
"rem": [["!limit", "x", max_length - 3, max_length]],
|
||||
},
|
||||
)
|
||||
+84
-21
@@ -2,8 +2,10 @@ import datetime
|
||||
import hashlib
|
||||
import logging
|
||||
import re
|
||||
from typing import Optional
|
||||
|
||||
from dateutil import parser
|
||||
from dateutil.relativedelta import relativedelta
|
||||
|
||||
from posthog.utils import convert_to_datetime_aware, is_valid_regex
|
||||
|
||||
@@ -11,6 +13,8 @@ __LONG_SCALE__ = float(0xFFFFFFFFFFFFFFF)
|
||||
|
||||
log = logging.getLogger("posthog")
|
||||
|
||||
NONE_VALUES_ALLOWED_OPERATORS = ["is_not"]
|
||||
|
||||
|
||||
class InconclusiveMatchError(Exception):
|
||||
pass
|
||||
@@ -117,15 +121,20 @@ def match_property(property, property_values) -> bool:
|
||||
|
||||
override_value = property_values[key]
|
||||
|
||||
if operator == "exact":
|
||||
if isinstance(value, list):
|
||||
return override_value in value
|
||||
return value == override_value
|
||||
if (operator not in NONE_VALUES_ALLOWED_OPERATORS) and override_value is None:
|
||||
return False
|
||||
|
||||
if operator == "is_not":
|
||||
if isinstance(value, list):
|
||||
return override_value not in value
|
||||
return value != override_value
|
||||
if operator in ("exact", "is_not"):
|
||||
|
||||
def compute_exact_match(value, override_value):
|
||||
if isinstance(value, list):
|
||||
return str(override_value).lower() in [str(val).lower() for val in value]
|
||||
return str(value).lower() == str(override_value).lower()
|
||||
|
||||
if operator == "exact":
|
||||
return compute_exact_match(value, override_value)
|
||||
else:
|
||||
return not compute_exact_match(value, override_value)
|
||||
|
||||
if operator == "is_set":
|
||||
return key in property_values
|
||||
@@ -142,23 +151,46 @@ def match_property(property, property_values) -> bool:
|
||||
if operator == "not_regex":
|
||||
return is_valid_regex(str(value)) and re.compile(str(value)).search(str(override_value)) is None
|
||||
|
||||
if operator == "gt":
|
||||
return type(override_value) == type(value) and override_value > value
|
||||
if operator in ("gt", "gte", "lt", "lte"):
|
||||
# :TRICKY: We adjust comparison based on the override value passed in,
|
||||
# to make sure we handle both numeric and string comparisons appropriately.
|
||||
def compare(lhs, rhs, operator):
|
||||
if operator == "gt":
|
||||
return lhs > rhs
|
||||
elif operator == "gte":
|
||||
return lhs >= rhs
|
||||
elif operator == "lt":
|
||||
return lhs < rhs
|
||||
elif operator == "lte":
|
||||
return lhs <= rhs
|
||||
else:
|
||||
raise ValueError(f"Invalid operator: {operator}")
|
||||
|
||||
if operator == "gte":
|
||||
return type(override_value) == type(value) and override_value >= value
|
||||
parsed_value = None
|
||||
try:
|
||||
parsed_value = float(value) # type: ignore
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if operator == "lt":
|
||||
return type(override_value) == type(value) and override_value < value
|
||||
|
||||
if operator == "lte":
|
||||
return type(override_value) == type(value) and override_value <= value
|
||||
if parsed_value is not None and override_value is not None:
|
||||
if isinstance(override_value, str):
|
||||
return compare(override_value, str(value), operator)
|
||||
else:
|
||||
return compare(override_value, parsed_value, operator)
|
||||
else:
|
||||
return compare(str(override_value), str(value), operator)
|
||||
|
||||
if operator in ["is_date_before", "is_date_after"]:
|
||||
try:
|
||||
parsed_date = parser.parse(value)
|
||||
parsed_date = convert_to_datetime_aware(parsed_date)
|
||||
except Exception:
|
||||
parsed_date = relative_date_parse_for_feature_flag_matching(str(value))
|
||||
|
||||
if not parsed_date:
|
||||
parsed_date = parser.parse(str(value))
|
||||
parsed_date = convert_to_datetime_aware(parsed_date)
|
||||
except Exception as e:
|
||||
raise InconclusiveMatchError("The date set on the flag is not a valid format") from e
|
||||
|
||||
if not parsed_date:
|
||||
raise InconclusiveMatchError("The date set on the flag is not a valid format")
|
||||
|
||||
if isinstance(override_value, datetime.datetime):
|
||||
@@ -185,7 +217,8 @@ def match_property(property, property_values) -> bool:
|
||||
else:
|
||||
raise InconclusiveMatchError("The date provided must be a string or date object")
|
||||
|
||||
return False
|
||||
# if we get here, we don't know how to handle the operator
|
||||
raise InconclusiveMatchError(f"Unknown operator {operator}")
|
||||
|
||||
|
||||
def match_cohort(property, property_values, cohort_properties) -> bool:
|
||||
@@ -271,3 +304,33 @@ def match_property_group(property_group, property_values, cohort_properties) ->
|
||||
|
||||
# if we get here, all matched in AND case, or none matched in OR case
|
||||
return property_group_type == "AND"
|
||||
|
||||
|
||||
def relative_date_parse_for_feature_flag_matching(value: str) -> Optional[datetime.datetime]:
|
||||
regex = r"^-?(?P<number>[0-9]+)(?P<interval>[a-z])$"
|
||||
match = re.search(regex, value)
|
||||
parsed_dt = datetime.datetime.now(datetime.timezone.utc)
|
||||
if match:
|
||||
number = int(match.group("number"))
|
||||
|
||||
if number >= 10_000:
|
||||
# Guard against overflow, disallow numbers greater than 10_000
|
||||
return None
|
||||
|
||||
interval = match.group("interval")
|
||||
if interval == "h":
|
||||
parsed_dt = parsed_dt - relativedelta(hours=number)
|
||||
elif interval == "d":
|
||||
parsed_dt = parsed_dt - relativedelta(days=number)
|
||||
elif interval == "w":
|
||||
parsed_dt = parsed_dt - relativedelta(weeks=number)
|
||||
elif interval == "m":
|
||||
parsed_dt = parsed_dt - relativedelta(months=number)
|
||||
elif interval == "y":
|
||||
parsed_dt = parsed_dt - relativedelta(years=number)
|
||||
else:
|
||||
return None
|
||||
|
||||
return parsed_dt
|
||||
else:
|
||||
return None
|
||||
|
||||
+16
-2
@@ -13,17 +13,31 @@ from posthog.version import VERSION
|
||||
|
||||
_session = requests.sessions.Session()
|
||||
|
||||
DEFAULT_HOST = "https://app.posthog.com"
|
||||
US_INGESTION_ENDPOINT = "https://us.i.posthog.com"
|
||||
EU_INGESTION_ENDPOINT = "https://eu.i.posthog.com"
|
||||
DEFAULT_HOST = US_INGESTION_ENDPOINT
|
||||
USER_AGENT = "posthog-python/" + VERSION
|
||||
|
||||
|
||||
def determine_server_host(host: Optional[str]) -> str:
|
||||
"""Determines the server host to use."""
|
||||
host_or_default = host or DEFAULT_HOST
|
||||
trimmed_host = remove_trailing_slash(host_or_default)
|
||||
if trimmed_host in ("https://app.posthog.com", "https://us.posthog.com"):
|
||||
return US_INGESTION_ENDPOINT
|
||||
elif trimmed_host == "https://eu.posthog.com":
|
||||
return EU_INGESTION_ENDPOINT
|
||||
else:
|
||||
return host_or_default
|
||||
|
||||
|
||||
def post(
|
||||
api_key: str, host: Optional[str] = None, path=None, gzip: bool = False, timeout: int = 15, **kwargs
|
||||
) -> requests.Response:
|
||||
"""Post the `kwargs` to the API"""
|
||||
log = logging.getLogger("posthog")
|
||||
body = kwargs
|
||||
body["sentAt"] = datetime.utcnow().replace(tzinfo=tzutc()).isoformat()
|
||||
body["sentAt"] = datetime.now(tz=tzutc()).isoformat()
|
||||
url = remove_trailing_slash(host or DEFAULT_HOST) + path
|
||||
body["api_key"] = api_key
|
||||
data = json.dumps(body, cls=DatetimeSerializer)
|
||||
|
||||
@@ -43,9 +43,9 @@ class PostHogIntegration(Integration):
|
||||
not not Hub.current.client.dsn and Dsn(Hub.current.client.dsn).project_id
|
||||
)
|
||||
if project_id:
|
||||
properties[
|
||||
"$sentry_url"
|
||||
] = f"{PostHogIntegration.prefix}{PostHogIntegration.organization}/issues/?project={project_id}&query={event['event_id']}"
|
||||
properties["$sentry_url"] = (
|
||||
f"{PostHogIntegration.prefix}{PostHogIntegration.organization}/issues/?project={project_id}&query={event['event_id']}"
|
||||
)
|
||||
|
||||
posthog.capture(posthog_distinct_id, "$exception", properties)
|
||||
|
||||
|
||||
@@ -0,0 +1,327 @@
|
||||
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"},
|
||||
]
|
||||
@@ -0,0 +1,5 @@
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("langchain")
|
||||
pytest.importorskip("langchain_community")
|
||||
pytest.importorskip("langgraph")
|
||||
@@ -0,0 +1,983 @@
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
from typing import List, Optional, TypedDict, Union
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from langchain_anthropic.chat_models import ChatAnthropic
|
||||
from langchain_community.chat_models.fake import FakeMessagesListChatModel
|
||||
from langchain_community.llms.fake import FakeListLLM, FakeStreamingListLLM
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langchain_core.prompts import ChatPromptTemplate
|
||||
from langchain_core.runnables import RunnableLambda
|
||||
from langchain_openai.chat_models import ChatOpenAI
|
||||
from langgraph.graph.state import END, START, StateGraph
|
||||
|
||||
from posthog.ai.langchain import CallbackHandler
|
||||
|
||||
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
|
||||
ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY")
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def mock_client():
|
||||
with patch("posthog.client.Client") as mock_client:
|
||||
mock_client.privacy_mode = False
|
||||
logging.getLogger("posthog").setLevel(logging.DEBUG)
|
||||
yield mock_client
|
||||
|
||||
|
||||
def test_parent_capture(mock_client):
|
||||
callbacks = CallbackHandler(mock_client)
|
||||
parent_run_id = uuid.uuid4()
|
||||
run_id = uuid.uuid4()
|
||||
callbacks._set_parent_of_run(run_id, parent_run_id)
|
||||
assert callbacks._parent_tree == {run_id: parent_run_id}
|
||||
callbacks._pop_parent_of_run(run_id)
|
||||
assert callbacks._parent_tree == {}
|
||||
callbacks._pop_parent_of_run(parent_run_id) # should not raise
|
||||
|
||||
|
||||
def test_find_root_run(mock_client):
|
||||
callbacks = CallbackHandler(mock_client)
|
||||
root_run_id = uuid.uuid4()
|
||||
parent_run_id = uuid.uuid4()
|
||||
run_id = uuid.uuid4()
|
||||
callbacks._set_parent_of_run(run_id, parent_run_id)
|
||||
callbacks._set_parent_of_run(parent_run_id, root_run_id)
|
||||
assert callbacks._find_root_run(run_id) == root_run_id
|
||||
new_run_id = uuid.uuid4()
|
||||
assert callbacks._find_root_run(new_run_id) == new_run_id
|
||||
|
||||
|
||||
def test_trace_id_generation(mock_client):
|
||||
callbacks = CallbackHandler(mock_client)
|
||||
run_id = uuid.uuid4()
|
||||
with patch("uuid.uuid4", return_value=run_id):
|
||||
assert callbacks._get_trace_id(run_id) == run_id
|
||||
run_id = uuid.uuid4()
|
||||
callbacks = CallbackHandler(mock_client, trace_id=run_id)
|
||||
assert callbacks._get_trace_id(uuid.uuid4()) == run_id
|
||||
|
||||
|
||||
def test_metadata_capture(mock_client):
|
||||
callbacks = CallbackHandler(mock_client)
|
||||
run_id = uuid.uuid4()
|
||||
with patch("time.time", return_value=1234567890):
|
||||
callbacks._set_run_metadata(
|
||||
{"kwargs": {"openai_api_base": "https://us.posthog.com"}},
|
||||
run_id,
|
||||
messages=[{"role": "user", "content": "Who won the world series in 2020?"}],
|
||||
invocation_params={"temperature": 0.5},
|
||||
metadata={"ls_model_name": "hog-mini", "ls_provider": "posthog"},
|
||||
)
|
||||
expected = {
|
||||
"model": "hog-mini",
|
||||
"messages": [{"role": "user", "content": "Who won the world series in 2020?"}],
|
||||
"start_time": 1234567890,
|
||||
"model_params": {"temperature": 0.5},
|
||||
"provider": "posthog",
|
||||
"base_url": "https://us.posthog.com",
|
||||
}
|
||||
assert callbacks._runs[run_id] == expected
|
||||
with patch("time.time", return_value=1234567891):
|
||||
run = callbacks._pop_run_metadata(run_id)
|
||||
assert run == {**expected, "end_time": 1234567891}
|
||||
assert callbacks._runs == {}
|
||||
callbacks._pop_run_metadata(uuid.uuid4()) # should not raise
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream", [True, False])
|
||||
def test_basic_chat_chain(mock_client, stream):
|
||||
prompt = ChatPromptTemplate.from_messages(
|
||||
[
|
||||
("system", "You are a helpful assistant."),
|
||||
("user", "Who won the world series in 2020?"),
|
||||
]
|
||||
)
|
||||
model = FakeMessagesListChatModel(
|
||||
responses=[
|
||||
AIMessage(
|
||||
content="The Los Angeles Dodgers won the World Series in 2020.",
|
||||
usage_metadata={
|
||||
"input_tokens": 10,
|
||||
"output_tokens": 10,
|
||||
"total_tokens": 20,
|
||||
},
|
||||
)
|
||||
]
|
||||
)
|
||||
callbacks = [CallbackHandler(mock_client)]
|
||||
chain = prompt | model
|
||||
if stream:
|
||||
result = [m for m in chain.stream({}, config={"callbacks": callbacks})][0]
|
||||
else:
|
||||
result = chain.invoke({}, config={"callbacks": callbacks})
|
||||
|
||||
assert result.content == "The Los Angeles Dodgers won the World Series in 2020."
|
||||
assert mock_client.capture.call_count == 2
|
||||
generation_args = mock_client.capture.call_args_list[0][1]
|
||||
generation_props = generation_args["properties"]
|
||||
trace_args = mock_client.capture.call_args_list[1][1]
|
||||
|
||||
assert generation_args["event"] == "$ai_generation"
|
||||
assert "distinct_id" in generation_args
|
||||
assert "$ai_model" in generation_props
|
||||
assert "$ai_provider" in generation_props
|
||||
assert generation_props["$ai_input"] == [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": "Who won the world series in 2020?"},
|
||||
]
|
||||
assert generation_props["$ai_output_choices"] == [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "The Los Angeles Dodgers won the World Series in 2020.",
|
||||
}
|
||||
]
|
||||
assert generation_props["$ai_input_tokens"] == 10
|
||||
assert generation_props["$ai_output_tokens"] == 10
|
||||
assert generation_props["$ai_http_status"] == 200
|
||||
assert generation_props["$ai_trace_id"] is not None
|
||||
assert isinstance(generation_props["$ai_latency"], float)
|
||||
assert trace_args["event"] == "$ai_trace"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream", [True, False])
|
||||
async def test_async_basic_chat_chain(mock_client, stream):
|
||||
prompt = ChatPromptTemplate.from_messages(
|
||||
[
|
||||
("system", "You are a helpful assistant."),
|
||||
("user", "Who won the world series in 2020?"),
|
||||
]
|
||||
)
|
||||
model = FakeMessagesListChatModel(
|
||||
responses=[
|
||||
AIMessage(
|
||||
content="The Los Angeles Dodgers won the World Series in 2020.",
|
||||
usage_metadata={
|
||||
"input_tokens": 10,
|
||||
"output_tokens": 10,
|
||||
"total_tokens": 20,
|
||||
},
|
||||
)
|
||||
]
|
||||
)
|
||||
callbacks = [CallbackHandler(mock_client)]
|
||||
chain = prompt | model
|
||||
if stream:
|
||||
result = [m async for m in chain.astream({}, config={"callbacks": callbacks})][0]
|
||||
else:
|
||||
result = await chain.ainvoke({}, config={"callbacks": callbacks})
|
||||
assert result.content == "The Los Angeles Dodgers won the World Series in 2020."
|
||||
assert mock_client.capture.call_count == 2
|
||||
|
||||
generation_args = mock_client.capture.call_args_list[0][1]
|
||||
generation_props = generation_args["properties"]
|
||||
trace_args = mock_client.capture.call_args_list[1][1]
|
||||
trace_props = trace_args["properties"]
|
||||
|
||||
assert generation_args["event"] == "$ai_generation"
|
||||
assert "distinct_id" in generation_args
|
||||
assert "$ai_model" in generation_props
|
||||
assert "$ai_provider" in generation_props
|
||||
assert generation_props["$ai_input"] == [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": "Who won the world series in 2020?"},
|
||||
]
|
||||
assert generation_props["$ai_output_choices"] == [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "The Los Angeles Dodgers won the World Series in 2020.",
|
||||
}
|
||||
]
|
||||
assert generation_props["$ai_input_tokens"] == 10
|
||||
assert generation_props["$ai_output_tokens"] == 10
|
||||
assert generation_props["$ai_http_status"] == 200
|
||||
assert generation_props["$ai_trace_id"] is not None
|
||||
assert isinstance(generation_props["$ai_latency"], float)
|
||||
|
||||
assert trace_args["event"] == "$ai_trace"
|
||||
assert "distinct_id" in generation_args
|
||||
assert trace_props["$ai_trace_id"] == generation_props["$ai_trace_id"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"Model,stream",
|
||||
[
|
||||
(FakeListLLM, True),
|
||||
(FakeListLLM, False),
|
||||
(FakeStreamingListLLM, True),
|
||||
(FakeStreamingListLLM, False),
|
||||
],
|
||||
)
|
||||
def test_basic_llm_chain(mock_client, Model, stream):
|
||||
model = Model(responses=["The Los Angeles Dodgers won the World Series in 2020."])
|
||||
callbacks: List[CallbackHandler] = [CallbackHandler(mock_client)]
|
||||
|
||||
if stream:
|
||||
result = "".join(
|
||||
[m for m in model.stream("Who won the world series in 2020?", config={"callbacks": callbacks})]
|
||||
)
|
||||
else:
|
||||
result = model.invoke("Who won the world series in 2020?", config={"callbacks": callbacks})
|
||||
assert result == "The Los Angeles Dodgers won the World Series in 2020."
|
||||
|
||||
assert mock_client.capture.call_count == 1
|
||||
args = mock_client.capture.call_args_list[0][1]
|
||||
props = args["properties"]
|
||||
|
||||
assert args["event"] == "$ai_generation"
|
||||
assert "distinct_id" in args
|
||||
assert "$ai_model" in props
|
||||
assert "$ai_provider" in props
|
||||
assert props["$ai_input"] == ["Who won the world series in 2020?"]
|
||||
assert props["$ai_output_choices"] == ["The Los Angeles Dodgers won the World Series in 2020."]
|
||||
assert props["$ai_http_status"] == 200
|
||||
assert props["$ai_trace_id"] is not None
|
||||
assert isinstance(props["$ai_latency"], float)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"Model,stream",
|
||||
[
|
||||
(FakeListLLM, True),
|
||||
(FakeListLLM, False),
|
||||
(FakeStreamingListLLM, True),
|
||||
(FakeStreamingListLLM, False),
|
||||
],
|
||||
)
|
||||
async def test_async_basic_llm_chain(mock_client, Model, stream):
|
||||
model = Model(responses=["The Los Angeles Dodgers won the World Series in 2020."])
|
||||
callbacks: List[CallbackHandler] = [CallbackHandler(mock_client)]
|
||||
|
||||
if stream:
|
||||
result = "".join(
|
||||
[m async for m in model.astream("Who won the world series in 2020?", config={"callbacks": callbacks})]
|
||||
)
|
||||
else:
|
||||
result = await model.ainvoke("Who won the world series in 2020?", config={"callbacks": callbacks})
|
||||
assert result == "The Los Angeles Dodgers won the World Series in 2020."
|
||||
|
||||
assert mock_client.capture.call_count == 1
|
||||
args = mock_client.capture.call_args_list[0][1]
|
||||
props = args["properties"]
|
||||
|
||||
assert args["event"] == "$ai_generation"
|
||||
assert "distinct_id" in args
|
||||
assert "$ai_model" in props
|
||||
assert "$ai_provider" in props
|
||||
assert props["$ai_input"] == ["Who won the world series in 2020?"]
|
||||
assert props["$ai_output_choices"] == ["The Los Angeles Dodgers won the World Series in 2020."]
|
||||
assert props["$ai_http_status"] == 200
|
||||
assert props["$ai_trace_id"] is not None
|
||||
assert isinstance(props["$ai_latency"], float)
|
||||
|
||||
|
||||
def test_trace_id_for_multiple_chains(mock_client):
|
||||
prompt = ChatPromptTemplate.from_messages(
|
||||
[
|
||||
("user", "Foo"),
|
||||
]
|
||||
)
|
||||
model = FakeMessagesListChatModel(responses=[AIMessage(content="Bar")])
|
||||
callbacks = [CallbackHandler(mock_client)]
|
||||
chain = prompt | model | RunnableLambda(lambda x: [x]) | model
|
||||
result = chain.invoke({}, config={"callbacks": callbacks})
|
||||
|
||||
assert result.content == "Bar"
|
||||
assert mock_client.capture.call_count == 3
|
||||
|
||||
first_call_args = mock_client.capture.call_args_list[0][1]
|
||||
first_call_props = first_call_args["properties"]
|
||||
assert first_call_args["event"] == "$ai_generation"
|
||||
assert "distinct_id" in first_call_args
|
||||
assert "$ai_model" in first_call_props
|
||||
assert "$ai_provider" in first_call_props
|
||||
assert first_call_props["$ai_input"] == [{"role": "user", "content": "Foo"}]
|
||||
assert first_call_props["$ai_output_choices"] == [{"role": "assistant", "content": "Bar"}]
|
||||
assert first_call_props["$ai_http_status"] == 200
|
||||
assert first_call_props["$ai_trace_id"] is not None
|
||||
assert isinstance(first_call_props["$ai_latency"], float)
|
||||
|
||||
second_generation_args = mock_client.capture.call_args_list[1][1]
|
||||
second_generation_props = second_generation_args["properties"]
|
||||
assert second_generation_args["event"] == "$ai_generation"
|
||||
assert "distinct_id" in second_generation_args
|
||||
assert "$ai_model" in second_generation_props
|
||||
assert "$ai_provider" in second_generation_props
|
||||
assert second_generation_props["$ai_input"] == [{"role": "assistant", "content": "Bar"}]
|
||||
assert second_generation_props["$ai_output_choices"] == [{"role": "assistant", "content": "Bar"}]
|
||||
assert second_generation_props["$ai_http_status"] == 200
|
||||
assert second_generation_props["$ai_trace_id"] is not None
|
||||
assert isinstance(second_generation_props["$ai_latency"], float)
|
||||
|
||||
trace_args = mock_client.capture.call_args_list[2][1]
|
||||
trace_props = trace_args["properties"]
|
||||
assert trace_args["event"] == "$ai_trace"
|
||||
assert "distinct_id" in trace_args
|
||||
assert trace_props["$ai_input_state"] == {}
|
||||
assert isinstance(trace_props["$ai_output_state"], AIMessage)
|
||||
assert trace_props["$ai_output_state"].content == "Bar"
|
||||
assert trace_props["$ai_trace_id"] is not None
|
||||
assert trace_props["$ai_trace_name"] == "RunnableSequence"
|
||||
|
||||
# Check that the trace_id is the same as the first call
|
||||
assert first_call_props["$ai_trace_id"] == second_generation_props["$ai_trace_id"]
|
||||
assert first_call_props["$ai_trace_id"] == trace_props["$ai_trace_id"]
|
||||
|
||||
|
||||
def test_personless_mode(mock_client):
|
||||
prompt = ChatPromptTemplate.from_messages([("user", "Foo")])
|
||||
chain = prompt | FakeMessagesListChatModel(responses=[AIMessage(content="Bar")])
|
||||
chain.invoke({}, config={"callbacks": [CallbackHandler(mock_client)]})
|
||||
assert mock_client.capture.call_count == 2
|
||||
generation_args = mock_client.capture.call_args_list[0][1]
|
||||
trace_args = mock_client.capture.call_args_list[1][1]
|
||||
assert generation_args["event"] == "$ai_generation"
|
||||
assert generation_args["properties"]["$process_person_profile"] is False
|
||||
assert trace_args["event"] == "$ai_trace"
|
||||
assert trace_args["properties"]["$process_person_profile"] is False
|
||||
|
||||
id = uuid.uuid4()
|
||||
chain.invoke({}, config={"callbacks": [CallbackHandler(mock_client, distinct_id=id)]})
|
||||
assert mock_client.capture.call_count == 4
|
||||
generation_args = mock_client.capture.call_args_list[2][1]
|
||||
trace_args = mock_client.capture.call_args_list[3][1]
|
||||
assert "$process_person_profile" not in generation_args["properties"]
|
||||
assert generation_args["distinct_id"] == id
|
||||
assert "$process_person_profile" not in trace_args["properties"]
|
||||
assert trace_args["distinct_id"] == id
|
||||
|
||||
|
||||
def test_personless_mode_exception(mock_client):
|
||||
prompt = ChatPromptTemplate.from_messages([("user", "Foo")])
|
||||
chain = prompt | ChatOpenAI(api_key="test", model="gpt-4o-mini")
|
||||
callbacks = CallbackHandler(mock_client)
|
||||
with pytest.raises(Exception):
|
||||
chain.invoke({}, config={"callbacks": [callbacks]})
|
||||
assert mock_client.capture.call_count == 2
|
||||
generation_args = mock_client.capture.call_args_list[0][1]
|
||||
trace_args = mock_client.capture.call_args_list[1][1]
|
||||
assert generation_args["event"] == "$ai_generation"
|
||||
assert generation_args["properties"]["$process_person_profile"] is False
|
||||
assert trace_args["event"] == "$ai_trace"
|
||||
assert trace_args["properties"]["$process_person_profile"] is False
|
||||
|
||||
id = uuid.uuid4()
|
||||
with pytest.raises(Exception):
|
||||
chain.invoke({}, config={"callbacks": [CallbackHandler(mock_client, distinct_id=id)]})
|
||||
assert mock_client.capture.call_count == 4
|
||||
generation_args = mock_client.capture.call_args_list[2][1]
|
||||
trace_args = mock_client.capture.call_args_list[3][1]
|
||||
assert "$process_person_profile" not in generation_args["properties"]
|
||||
assert generation_args["distinct_id"] == id
|
||||
assert "$process_person_profile" not in trace_args["properties"]
|
||||
assert trace_args["distinct_id"] == id
|
||||
|
||||
|
||||
def test_metadata(mock_client):
|
||||
prompt = ChatPromptTemplate.from_messages(
|
||||
[
|
||||
("user", "Foo"),
|
||||
]
|
||||
)
|
||||
model = FakeMessagesListChatModel(responses=[AIMessage(content="Bar")])
|
||||
callbacks = [
|
||||
CallbackHandler(
|
||||
mock_client,
|
||||
trace_id="test-trace-id",
|
||||
distinct_id="test_id",
|
||||
properties={"foo": "bar"},
|
||||
)
|
||||
]
|
||||
chain = prompt | model
|
||||
result = chain.invoke({"plan": None}, config={"callbacks": callbacks})
|
||||
|
||||
assert result.content == "Bar"
|
||||
assert mock_client.capture.call_count == 2
|
||||
|
||||
generation_call_args = mock_client.capture.call_args_list[0][1]
|
||||
generation_call_props = generation_call_args["properties"]
|
||||
assert generation_call_args["distinct_id"] == "test_id"
|
||||
assert generation_call_args["event"] == "$ai_generation"
|
||||
assert generation_call_props["$ai_trace_id"] == "test-trace-id"
|
||||
assert generation_call_props["foo"] == "bar"
|
||||
assert generation_call_props["$ai_input"] == [{"role": "user", "content": "Foo"}]
|
||||
assert generation_call_props["$ai_output_choices"] == [{"role": "assistant", "content": "Bar"}]
|
||||
assert generation_call_props["$ai_http_status"] == 200
|
||||
assert isinstance(generation_call_props["$ai_latency"], float)
|
||||
|
||||
trace_call_args = mock_client.capture.call_args_list[1][1]
|
||||
trace_call_props = trace_call_args["properties"]
|
||||
assert trace_call_args["distinct_id"] == "test_id"
|
||||
assert trace_call_args["event"] == "$ai_trace"
|
||||
assert trace_call_props["$ai_trace_id"] == "test-trace-id"
|
||||
assert trace_call_props["$ai_trace_name"] == "RunnableSequence"
|
||||
assert trace_call_props["foo"] == "bar"
|
||||
assert trace_call_props["$ai_input_state"] == {"plan": None}
|
||||
assert isinstance(trace_call_props["$ai_output_state"], AIMessage)
|
||||
assert trace_call_props["$ai_output_state"].content == "Bar"
|
||||
|
||||
|
||||
class FakeGraphState(TypedDict):
|
||||
messages: List[Union[HumanMessage, AIMessage]]
|
||||
xyz: Optional[str]
|
||||
|
||||
|
||||
def test_graph_state(mock_client):
|
||||
config = {"callbacks": [CallbackHandler(mock_client)]}
|
||||
|
||||
graph = StateGraph(FakeGraphState)
|
||||
graph.add_node(
|
||||
"fake_plain",
|
||||
lambda state: {
|
||||
"messages": [
|
||||
*state["messages"],
|
||||
AIMessage(content="Let's explore bar."),
|
||||
],
|
||||
"xyz": "abc",
|
||||
},
|
||||
)
|
||||
intermediate_chain = ChatPromptTemplate.from_messages(
|
||||
[("user", "Question: What's a bar?")]
|
||||
) | FakeMessagesListChatModel(
|
||||
responses=[
|
||||
AIMessage(content="It's a type of greeble."),
|
||||
]
|
||||
)
|
||||
graph.add_node(
|
||||
"fake_llm",
|
||||
lambda state: {
|
||||
"messages": [
|
||||
*state["messages"],
|
||||
intermediate_chain.invoke(state),
|
||||
],
|
||||
"xyz": state["xyz"],
|
||||
},
|
||||
)
|
||||
graph.add_edge(START, "fake_plain")
|
||||
graph.add_edge("fake_plain", "fake_llm")
|
||||
graph.add_edge("fake_llm", END)
|
||||
|
||||
result = graph.compile().invoke(
|
||||
{"messages": [HumanMessage(content="What's a bar?")], "xyz": None},
|
||||
config=config,
|
||||
)
|
||||
|
||||
assert len(result["messages"]) == 3
|
||||
assert isinstance(result["messages"][0], HumanMessage)
|
||||
assert result["messages"][0].content == "What's a bar?"
|
||||
assert isinstance(result["messages"][1], AIMessage)
|
||||
assert result["messages"][1].content == "Let's explore bar."
|
||||
assert isinstance(result["messages"][2], AIMessage)
|
||||
assert result["messages"][2].content == "It's a type of greeble."
|
||||
|
||||
assert mock_client.capture.call_count == 2
|
||||
generation_args = mock_client.capture.call_args_list[0][1]
|
||||
trace_args = mock_client.capture.call_args_list[1][1]
|
||||
assert generation_args["event"] == "$ai_generation"
|
||||
assert trace_args["event"] == "$ai_trace"
|
||||
assert trace_args["properties"]["$ai_trace_name"] == "LangGraph"
|
||||
|
||||
assert len(trace_args["properties"]["$ai_input_state"]["messages"]) == 1
|
||||
assert isinstance(trace_args["properties"]["$ai_input_state"]["messages"][0], HumanMessage)
|
||||
assert trace_args["properties"]["$ai_input_state"]["messages"][0].content == "What's a bar?"
|
||||
assert trace_args["properties"]["$ai_input_state"]["messages"][0].type == "human"
|
||||
assert trace_args["properties"]["$ai_input_state"]["xyz"] is None
|
||||
assert len(trace_args["properties"]["$ai_output_state"]["messages"]) == 3
|
||||
|
||||
assert isinstance(trace_args["properties"]["$ai_output_state"]["messages"][0], HumanMessage)
|
||||
assert trace_args["properties"]["$ai_output_state"]["messages"][0].content == "What's a bar?"
|
||||
assert isinstance(trace_args["properties"]["$ai_output_state"]["messages"][1], AIMessage)
|
||||
assert trace_args["properties"]["$ai_output_state"]["messages"][1].content == "Let's explore bar."
|
||||
assert isinstance(trace_args["properties"]["$ai_output_state"]["messages"][2], AIMessage)
|
||||
assert trace_args["properties"]["$ai_output_state"]["messages"][2].content == "It's a type of greeble."
|
||||
assert trace_args["properties"]["$ai_output_state"]["xyz"] == "abc"
|
||||
|
||||
|
||||
def test_callbacks_logic(mock_client):
|
||||
prompt = ChatPromptTemplate.from_messages([("user", "Foo")])
|
||||
model = FakeMessagesListChatModel(responses=[AIMessage(content="Bar")])
|
||||
callbacks = CallbackHandler(
|
||||
mock_client,
|
||||
trace_id="test-trace-id",
|
||||
distinct_id="test_id",
|
||||
properties={"foo": "bar"},
|
||||
)
|
||||
chain = prompt | model
|
||||
|
||||
chain.invoke({}, config={"callbacks": [callbacks]})
|
||||
assert callbacks._runs == {}
|
||||
assert callbacks._parent_tree == {}
|
||||
|
||||
def assert_intermediary_run(m):
|
||||
assert callbacks._runs == {}
|
||||
assert len(callbacks._parent_tree.items()) == 1
|
||||
return [m]
|
||||
|
||||
(chain | RunnableLambda(assert_intermediary_run) | model).invoke({}, config={"callbacks": [callbacks]})
|
||||
assert callbacks._runs == {}
|
||||
assert callbacks._parent_tree == {}
|
||||
|
||||
|
||||
def test_exception_in_chain(mock_client):
|
||||
def runnable(_):
|
||||
raise ValueError("test")
|
||||
|
||||
callbacks = CallbackHandler(mock_client)
|
||||
with pytest.raises(ValueError):
|
||||
RunnableLambda(runnable).invoke({}, config={"callbacks": [callbacks]})
|
||||
|
||||
assert callbacks._runs == {}
|
||||
assert callbacks._parent_tree == {}
|
||||
assert mock_client.capture.call_count == 1
|
||||
trace_call_args = mock_client.capture.call_args_list[0][1]
|
||||
assert trace_call_args["event"] == "$ai_trace"
|
||||
assert trace_call_args["properties"]["$ai_trace_name"] == "runnable"
|
||||
|
||||
|
||||
def test_openai_error(mock_client):
|
||||
prompt = ChatPromptTemplate.from_messages([("user", "Foo")])
|
||||
chain = prompt | ChatOpenAI(api_key="test", model="gpt-4o-mini")
|
||||
callbacks = CallbackHandler(mock_client)
|
||||
|
||||
# 401
|
||||
with pytest.raises(Exception):
|
||||
chain.invoke({}, config={"callbacks": [callbacks]})
|
||||
|
||||
assert callbacks._runs == {}
|
||||
assert callbacks._parent_tree == {}
|
||||
assert mock_client.capture.call_count == 2
|
||||
generation_args = mock_client.capture.call_args_list[0][1]
|
||||
props = generation_args["properties"]
|
||||
assert props["$ai_http_status"] == 401
|
||||
assert props["$ai_input"] == [{"role": "user", "content": "Foo"}]
|
||||
assert "$ai_output_choices" not in props
|
||||
|
||||
|
||||
@pytest.mark.skipif(not OPENAI_API_KEY, reason="OpenAI API key not set")
|
||||
def test_openai_chain(mock_client):
|
||||
prompt = ChatPromptTemplate.from_messages(
|
||||
[
|
||||
("system", 'You must always answer with "Bar".'),
|
||||
("user", "Foo"),
|
||||
]
|
||||
)
|
||||
chain = prompt | ChatOpenAI(
|
||||
api_key=OPENAI_API_KEY,
|
||||
model="gpt-4o-mini",
|
||||
temperature=0,
|
||||
max_tokens=1,
|
||||
)
|
||||
callbacks = CallbackHandler(
|
||||
mock_client,
|
||||
trace_id="test-trace-id",
|
||||
distinct_id="test_id",
|
||||
properties={"foo": "bar"},
|
||||
)
|
||||
start_time = time.time()
|
||||
result = chain.invoke({}, config={"callbacks": [callbacks]})
|
||||
approximate_latency = math.floor(time.time() - start_time)
|
||||
|
||||
assert result.content == "Bar"
|
||||
assert mock_client.capture.call_count == 2
|
||||
|
||||
first_call_args = mock_client.capture.call_args_list[0][1]
|
||||
first_call_props = first_call_args["properties"]
|
||||
assert first_call_args["event"] == "$ai_generation"
|
||||
assert first_call_props["$ai_trace_id"] == "test-trace-id"
|
||||
assert first_call_props["$ai_provider"] == "openai"
|
||||
assert first_call_props["$ai_model"] == "gpt-4o-mini"
|
||||
assert first_call_props["foo"] == "bar"
|
||||
|
||||
# langchain-openai for langchain v3
|
||||
if "max_completion_tokens" in first_call_props["$ai_model_parameters"]:
|
||||
assert first_call_props["$ai_model_parameters"] == {
|
||||
"temperature": 0.0,
|
||||
"max_completion_tokens": 1,
|
||||
"stream": False,
|
||||
}
|
||||
else:
|
||||
assert first_call_props["$ai_model_parameters"] == {
|
||||
"temperature": 0.0,
|
||||
"max_tokens": 1,
|
||||
"n": 1,
|
||||
"stream": False,
|
||||
}
|
||||
assert first_call_props["$ai_input"] == [
|
||||
{"role": "system", "content": 'You must always answer with "Bar".'},
|
||||
{"role": "user", "content": "Foo"},
|
||||
]
|
||||
assert first_call_props["$ai_output_choices"] == [{"role": "assistant", "content": "Bar", "refusal": None}]
|
||||
assert first_call_props["$ai_http_status"] == 200
|
||||
assert isinstance(first_call_props["$ai_latency"], float)
|
||||
assert min(approximate_latency - 1, 0) <= math.floor(first_call_props["$ai_latency"]) <= approximate_latency
|
||||
assert first_call_props["$ai_input_tokens"] == 20
|
||||
assert first_call_props["$ai_output_tokens"] == 1
|
||||
|
||||
|
||||
@pytest.mark.skipif(not OPENAI_API_KEY, reason="OpenAI API key not set")
|
||||
def test_openai_captures_multiple_generations(mock_client):
|
||||
prompt = ChatPromptTemplate.from_messages(
|
||||
[
|
||||
("system", 'You must always answer with "Bar".'),
|
||||
("user", "Foo"),
|
||||
]
|
||||
)
|
||||
chain = prompt | ChatOpenAI(
|
||||
api_key=OPENAI_API_KEY,
|
||||
model="gpt-4o-mini",
|
||||
temperature=0,
|
||||
max_tokens=1,
|
||||
n=2,
|
||||
)
|
||||
callbacks = CallbackHandler(mock_client)
|
||||
result = chain.invoke({}, config={"callbacks": [callbacks]})
|
||||
|
||||
assert result.content == "Bar"
|
||||
assert mock_client.capture.call_count == 2
|
||||
|
||||
first_call_args = mock_client.capture.call_args_list[0][1]
|
||||
first_call_props = first_call_args["properties"]
|
||||
second_call_args = mock_client.capture.call_args_list[1][1]
|
||||
second_call_props = second_call_args["properties"]
|
||||
|
||||
assert first_call_args["event"] == "$ai_generation"
|
||||
assert first_call_props["$ai_input"] == [
|
||||
{"role": "system", "content": 'You must always answer with "Bar".'},
|
||||
{"role": "user", "content": "Foo"},
|
||||
]
|
||||
assert first_call_props["$ai_output_choices"] == [
|
||||
{"role": "assistant", "content": "Bar", "refusal": None},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Bar",
|
||||
},
|
||||
]
|
||||
|
||||
# langchain-openai for langchain v3
|
||||
if "max_completion_tokens" in first_call_props["$ai_model_parameters"]:
|
||||
assert first_call_props["$ai_model_parameters"] == {
|
||||
"temperature": 0.0,
|
||||
"max_completion_tokens": 1,
|
||||
"stream": False,
|
||||
"n": 2,
|
||||
}
|
||||
else:
|
||||
assert first_call_props["$ai_model_parameters"] == {
|
||||
"temperature": 0.0,
|
||||
"max_tokens": 1,
|
||||
"stream": False,
|
||||
"n": 2,
|
||||
}
|
||||
assert first_call_props["$ai_http_status"] == 200
|
||||
|
||||
assert second_call_args["event"] == "$ai_trace"
|
||||
assert second_call_props["$ai_input_state"] == {}
|
||||
assert isinstance(second_call_props["$ai_output_state"], AIMessage)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not OPENAI_API_KEY, reason="OpenAI API key not set")
|
||||
def test_openai_streaming(mock_client):
|
||||
prompt = ChatPromptTemplate.from_messages(
|
||||
[
|
||||
("system", 'You must always answer with "Bar".'),
|
||||
("user", "Foo"),
|
||||
]
|
||||
)
|
||||
chain = prompt | ChatOpenAI(
|
||||
api_key=OPENAI_API_KEY,
|
||||
model="gpt-4o-mini",
|
||||
temperature=0,
|
||||
max_tokens=1,
|
||||
stream=True,
|
||||
stream_usage=True,
|
||||
)
|
||||
callbacks = CallbackHandler(mock_client)
|
||||
result = [m for m in chain.stream({}, config={"callbacks": [callbacks]})]
|
||||
result = sum(result[1:], result[0])
|
||||
|
||||
assert result.content == "Bar"
|
||||
assert mock_client.capture.call_count == 2
|
||||
|
||||
first_call_args = mock_client.capture.call_args_list[0][1]
|
||||
first_call_props = first_call_args["properties"]
|
||||
second_call_args = mock_client.capture.call_args_list[1][1]
|
||||
second_call_props = second_call_args["properties"]
|
||||
|
||||
assert first_call_args["event"] == "$ai_generation"
|
||||
assert first_call_props["$ai_model_parameters"]["stream"]
|
||||
assert first_call_props["$ai_input"] == [
|
||||
{"role": "system", "content": 'You must always answer with "Bar".'},
|
||||
{"role": "user", "content": "Foo"},
|
||||
]
|
||||
assert first_call_props["$ai_output_choices"] == [{"role": "assistant", "content": "Bar"}]
|
||||
assert first_call_props["$ai_http_status"] == 200
|
||||
assert first_call_props["$ai_input_tokens"] == 20
|
||||
assert first_call_props["$ai_output_tokens"] == 1
|
||||
|
||||
assert second_call_args["event"] == "$ai_trace"
|
||||
assert second_call_props["$ai_input_state"] == {"input": ""}
|
||||
assert isinstance(second_call_props["$ai_output_state"], AIMessage)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not OPENAI_API_KEY, reason="OpenAI API key not set")
|
||||
async def test_async_openai_streaming(mock_client):
|
||||
prompt = ChatPromptTemplate.from_messages(
|
||||
[
|
||||
("system", 'You must always answer with "Bar".'),
|
||||
("user", "Foo"),
|
||||
]
|
||||
)
|
||||
chain = prompt | ChatOpenAI(
|
||||
api_key=OPENAI_API_KEY,
|
||||
model="gpt-4o-mini",
|
||||
temperature=0,
|
||||
max_tokens=1,
|
||||
stream=True,
|
||||
stream_usage=True,
|
||||
)
|
||||
callbacks = CallbackHandler(mock_client)
|
||||
result = [m async for m in chain.astream({}, config={"callbacks": [callbacks]})]
|
||||
result = sum(result[1:], result[0])
|
||||
|
||||
assert result.content == "Bar"
|
||||
assert mock_client.capture.call_count == 2
|
||||
|
||||
first_call_args = mock_client.capture.call_args_list[0][1]
|
||||
first_call_props = first_call_args["properties"]
|
||||
second_call_args = mock_client.capture.call_args_list[1][1]
|
||||
second_call_props = second_call_args["properties"]
|
||||
|
||||
assert first_call_args["event"] == "$ai_generation"
|
||||
assert first_call_props["$ai_model_parameters"]["stream"]
|
||||
assert first_call_props["$ai_input"] == [
|
||||
{"role": "system", "content": 'You must always answer with "Bar".'},
|
||||
{"role": "user", "content": "Foo"},
|
||||
]
|
||||
assert first_call_props["$ai_output_choices"] == [{"role": "assistant", "content": "Bar"}]
|
||||
assert first_call_props["$ai_http_status"] == 200
|
||||
assert first_call_props["$ai_input_tokens"] == 20
|
||||
assert first_call_props["$ai_output_tokens"] == 1
|
||||
|
||||
assert second_call_args["event"] == "$ai_trace"
|
||||
assert second_call_props["$ai_input_state"] == {"input": ""}
|
||||
assert isinstance(second_call_props["$ai_output_state"], AIMessage)
|
||||
|
||||
|
||||
def test_base_url_retrieval(mock_client):
|
||||
prompt = ChatPromptTemplate.from_messages([("user", "Foo")])
|
||||
chain = prompt | ChatOpenAI(
|
||||
api_key="test",
|
||||
model="posthog-mini",
|
||||
base_url="https://test.posthog.com",
|
||||
)
|
||||
callbacks = CallbackHandler(mock_client)
|
||||
with pytest.raises(Exception):
|
||||
chain.invoke({}, config={"callbacks": [callbacks]})
|
||||
|
||||
assert mock_client.capture.call_count == 2
|
||||
generation_call = mock_client.capture.call_args_list[0][1]
|
||||
assert generation_call["properties"]["$ai_base_url"] == "https://test.posthog.com"
|
||||
|
||||
|
||||
def test_groups(mock_client):
|
||||
prompt = ChatPromptTemplate.from_messages(
|
||||
[
|
||||
("system", 'You must always answer with "Bar".'),
|
||||
("user", "Foo"),
|
||||
]
|
||||
)
|
||||
model = FakeMessagesListChatModel(responses=[AIMessage(content="Bar")])
|
||||
chain = prompt | model
|
||||
callbacks = CallbackHandler(mock_client, groups={"company": "test_company"})
|
||||
chain.invoke({}, config={"callbacks": [callbacks]})
|
||||
|
||||
assert mock_client.capture.call_count == 2
|
||||
generation_call = mock_client.capture.call_args_list[0][1]
|
||||
assert generation_call["groups"] == {"company": "test_company"}
|
||||
|
||||
|
||||
def test_privacy_mode_local(mock_client):
|
||||
prompt = ChatPromptTemplate.from_messages(
|
||||
[
|
||||
("system", 'You must always answer with "Bar".'),
|
||||
("user", "Foo"),
|
||||
]
|
||||
)
|
||||
model = FakeMessagesListChatModel(responses=[AIMessage(content="Bar")])
|
||||
chain = prompt | model
|
||||
callbacks = CallbackHandler(mock_client, privacy_mode=True)
|
||||
chain.invoke({}, config={"callbacks": [callbacks]})
|
||||
|
||||
assert mock_client.capture.call_count == 2
|
||||
generation_call = mock_client.capture.call_args_list[0][1]
|
||||
assert generation_call["properties"]["$ai_input"] is None
|
||||
assert generation_call["properties"]["$ai_output_choices"] is None
|
||||
|
||||
|
||||
def test_privacy_mode_global(mock_client):
|
||||
mock_client.privacy_mode = True
|
||||
prompt = ChatPromptTemplate.from_messages(
|
||||
[
|
||||
("system", 'You must always answer with "Bar".'),
|
||||
("user", "Foo"),
|
||||
]
|
||||
)
|
||||
model = FakeMessagesListChatModel(responses=[AIMessage(content="Bar")])
|
||||
chain = prompt | model
|
||||
callbacks = CallbackHandler(mock_client)
|
||||
chain.invoke({}, config={"callbacks": [callbacks]})
|
||||
|
||||
assert mock_client.capture.call_count == 2
|
||||
generation_call = mock_client.capture.call_args_list[0][1]
|
||||
assert generation_call["properties"]["$ai_input"] is None
|
||||
assert generation_call["properties"]["$ai_output_choices"] is None
|
||||
|
||||
|
||||
@pytest.mark.skipif(not ANTHROPIC_API_KEY, reason="ANTHROPIC_API_KEY is not set")
|
||||
def test_anthropic_chain(mock_client):
|
||||
prompt = ChatPromptTemplate.from_messages(
|
||||
[
|
||||
("system", 'You must always answer with "Bar".'),
|
||||
("user", "Foo"),
|
||||
]
|
||||
)
|
||||
chain = prompt | ChatAnthropic(
|
||||
api_key=ANTHROPIC_API_KEY,
|
||||
model="claude-3-opus-20240229",
|
||||
temperature=0,
|
||||
max_tokens=1,
|
||||
)
|
||||
callbacks = CallbackHandler(
|
||||
mock_client,
|
||||
trace_id="test-trace-id",
|
||||
distinct_id="test_id",
|
||||
properties={"foo": "bar"},
|
||||
)
|
||||
start_time = time.time()
|
||||
result = chain.invoke({}, config={"callbacks": [callbacks]})
|
||||
approximate_latency = math.floor(time.time() - start_time)
|
||||
|
||||
assert result.content == "Bar"
|
||||
assert mock_client.capture.call_count == 2
|
||||
|
||||
first_call_args = mock_client.capture.call_args_list[0][1]
|
||||
first_call_props = first_call_args["properties"]
|
||||
second_call_args = mock_client.capture.call_args_list[1][1]
|
||||
second_call_props = second_call_args["properties"]
|
||||
|
||||
assert first_call_args["event"] == "$ai_generation"
|
||||
assert first_call_props["$ai_trace_id"] == "test-trace-id"
|
||||
assert first_call_props["$ai_provider"] == "anthropic"
|
||||
assert first_call_props["$ai_model"] == "claude-3-opus-20240229"
|
||||
assert first_call_props["foo"] == "bar"
|
||||
|
||||
assert first_call_props["$ai_model_parameters"] == {
|
||||
"temperature": 0.0,
|
||||
"max_tokens": 1,
|
||||
"streaming": False,
|
||||
}
|
||||
assert first_call_props["$ai_input"] == [
|
||||
{"role": "system", "content": 'You must always answer with "Bar".'},
|
||||
{"role": "user", "content": "Foo"},
|
||||
]
|
||||
assert first_call_props["$ai_output_choices"] == [{"role": "assistant", "content": "Bar"}]
|
||||
assert first_call_props["$ai_http_status"] == 200
|
||||
assert isinstance(first_call_props["$ai_latency"], float)
|
||||
assert min(approximate_latency - 1, 0) <= math.floor(first_call_props["$ai_latency"]) <= approximate_latency
|
||||
assert first_call_props["$ai_input_tokens"] == 17
|
||||
assert first_call_props["$ai_output_tokens"] == 1
|
||||
|
||||
assert second_call_args["event"] == "$ai_trace"
|
||||
assert second_call_props["$ai_input_state"] == {}
|
||||
assert isinstance(second_call_props["$ai_output_state"], AIMessage)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not ANTHROPIC_API_KEY, reason="ANTHROPIC_API_KEY is not set")
|
||||
async def test_async_anthropic_streaming(mock_client):
|
||||
prompt = ChatPromptTemplate.from_messages(
|
||||
[
|
||||
("system", 'You must always answer with "Bar".'),
|
||||
("user", "Foo"),
|
||||
]
|
||||
)
|
||||
chain = prompt | ChatAnthropic(
|
||||
api_key=ANTHROPIC_API_KEY,
|
||||
model="claude-3-opus-20240229",
|
||||
temperature=0,
|
||||
max_tokens=1,
|
||||
streaming=True,
|
||||
stream_usage=True,
|
||||
)
|
||||
callbacks = CallbackHandler(mock_client)
|
||||
result = [m async for m in chain.astream({}, config={"callbacks": [callbacks]})]
|
||||
result = sum(result[1:], result[0])
|
||||
|
||||
assert result.content == "Bar"
|
||||
assert mock_client.capture.call_count == 2
|
||||
|
||||
first_call_args = mock_client.capture.call_args_list[0][1]
|
||||
first_call_props = first_call_args["properties"]
|
||||
second_call_args = mock_client.capture.call_args_list[1][1]
|
||||
second_call_props = second_call_args["properties"]
|
||||
|
||||
assert first_call_args["event"] == "$ai_generation"
|
||||
assert first_call_props["$ai_model_parameters"]["streaming"]
|
||||
assert first_call_props["$ai_input"] == [
|
||||
{"role": "system", "content": 'You must always answer with "Bar".'},
|
||||
{"role": "user", "content": "Foo"},
|
||||
]
|
||||
assert first_call_props["$ai_output_choices"] == [{"role": "assistant", "content": "Bar"}]
|
||||
assert first_call_props["$ai_http_status"] == 200
|
||||
assert first_call_props["$ai_input_tokens"] == 17
|
||||
assert first_call_props["$ai_output_tokens"] is not None
|
||||
|
||||
assert second_call_args["event"] == "$ai_trace"
|
||||
assert second_call_props["$ai_input_state"] == {
|
||||
"input": "",
|
||||
}
|
||||
assert isinstance(second_call_props["$ai_output_state"], AIMessage)
|
||||
|
||||
|
||||
def test_tool_calls(mock_client):
|
||||
prompt = ChatPromptTemplate.from_messages([("user", "Foo")])
|
||||
model = FakeMessagesListChatModel(
|
||||
responses=[
|
||||
AIMessage(
|
||||
content="Bar",
|
||||
additional_kwargs={
|
||||
"tool_calls": [
|
||||
{
|
||||
"type": "function",
|
||||
"id": "123",
|
||||
"function": {
|
||||
"name": "test",
|
||||
"args": '{"a": 1}',
|
||||
},
|
||||
}
|
||||
]
|
||||
},
|
||||
)
|
||||
]
|
||||
)
|
||||
chain = prompt | model
|
||||
callbacks = CallbackHandler(mock_client)
|
||||
chain.invoke({}, config={"callbacks": [callbacks]})
|
||||
|
||||
assert mock_client.capture.call_count == 2
|
||||
generation_call = mock_client.capture.call_args_list[0][1]
|
||||
assert generation_call["properties"]["$ai_output_choices"][0]["tool_calls"] == [
|
||||
{
|
||||
"type": "function",
|
||||
"id": "123",
|
||||
"function": {
|
||||
"name": "test",
|
||||
"args": '{"a": 1}',
|
||||
},
|
||||
}
|
||||
]
|
||||
assert "additional_kwargs" not in generation_call["properties"]["$ai_output_choices"][0]
|
||||
@@ -0,0 +1,175 @@
|
||||
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
|
||||
@@ -0,0 +1,3 @@
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("django")
|
||||
@@ -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",
|
||||
}
|
||||
+623
-7
@@ -84,6 +84,184 @@ 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"}}
|
||||
@@ -106,14 +284,187 @@ class TestClient(unittest.TestCase):
|
||||
self.assertEqual(patch_decide.call_count, 1)
|
||||
|
||||
@mock.patch("posthog.client.decide")
|
||||
def test_get_active_feature_flags(self, patch_decide):
|
||||
patch_decide.return_value = {
|
||||
"featureFlags": {"beta-feature": "random-variant", "alpha-feature": True, "off-feature": False}
|
||||
}
|
||||
|
||||
def test_basic_capture_with_locally_evaluated_feature_flags(self, patch_decide):
|
||||
patch_decide.return_value = {"featureFlags": {"beta-feature": "random-variant"}}
|
||||
client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, personal_api_key=FAKE_TEST_API_KEY)
|
||||
variants = client._get_active_feature_variants("some_id", None, None, None)
|
||||
self.assertEqual(variants, {"beta-feature": "random-variant", "alpha-feature": True})
|
||||
|
||||
multivariate_flag = {
|
||||
"id": 1,
|
||||
"name": "Beta Feature",
|
||||
"key": "beta-feature-local",
|
||||
"is_simple_flag": False,
|
||||
"active": True,
|
||||
"rollout_percentage": 100,
|
||||
"filters": {
|
||||
"groups": [
|
||||
{
|
||||
"properties": [
|
||||
{"key": "email", "type": "person", "value": "test@posthog.com", "operator": "exact"}
|
||||
],
|
||||
"rollout_percentage": 100,
|
||||
},
|
||||
{
|
||||
"rollout_percentage": 50,
|
||||
},
|
||||
],
|
||||
"multivariate": {
|
||||
"variants": [
|
||||
{"key": "first-variant", "name": "First Variant", "rollout_percentage": 50},
|
||||
{"key": "second-variant", "name": "Second Variant", "rollout_percentage": 25},
|
||||
{"key": "third-variant", "name": "Third Variant", "rollout_percentage": 25},
|
||||
]
|
||||
},
|
||||
"payloads": {"first-variant": "some-payload", "third-variant": {"a": "json"}},
|
||||
},
|
||||
}
|
||||
basic_flag = {
|
||||
"id": 1,
|
||||
"name": "Beta Feature",
|
||||
"key": "person-flag",
|
||||
"is_simple_flag": True,
|
||||
"active": True,
|
||||
"filters": {
|
||||
"groups": [
|
||||
{
|
||||
"properties": [
|
||||
{
|
||||
"key": "region",
|
||||
"operator": "exact",
|
||||
"value": ["USA"],
|
||||
"type": "person",
|
||||
}
|
||||
],
|
||||
"rollout_percentage": 100,
|
||||
}
|
||||
],
|
||||
"payloads": {"true": 300},
|
||||
},
|
||||
}
|
||||
false_flag = {
|
||||
"id": 1,
|
||||
"name": "Beta Feature",
|
||||
"key": "false-flag",
|
||||
"is_simple_flag": True,
|
||||
"active": True,
|
||||
"filters": {
|
||||
"groups": [
|
||||
{
|
||||
"properties": [],
|
||||
"rollout_percentage": 0,
|
||||
}
|
||||
],
|
||||
"payloads": {"true": 300},
|
||||
},
|
||||
}
|
||||
client.feature_flags = [multivariate_flag, basic_flag, false_flag]
|
||||
|
||||
success, msg = client.capture("distinct_id", "python test event")
|
||||
client.flush()
|
||||
self.assertTrue(success)
|
||||
self.assertFalse(self.failed)
|
||||
|
||||
self.assertEqual(msg["event"], "python test event")
|
||||
self.assertTrue(isinstance(msg["timestamp"], str))
|
||||
self.assertIsNone(msg.get("uuid"))
|
||||
self.assertEqual(msg["distinct_id"], "distinct_id")
|
||||
self.assertEqual(msg["properties"]["$lib"], "posthog-python")
|
||||
self.assertEqual(msg["properties"]["$lib_version"], VERSION)
|
||||
self.assertEqual(msg["properties"]["$feature/beta-feature-local"], "third-variant")
|
||||
self.assertEqual(msg["properties"]["$feature/false-flag"], False)
|
||||
self.assertEqual(msg["properties"]["$active_feature_flags"], ["beta-feature-local"])
|
||||
assert "$feature/beta-feature" not in msg["properties"]
|
||||
|
||||
self.assertEqual(patch_decide.call_count, 0)
|
||||
|
||||
# test that flags are not evaluated without local evaluation
|
||||
client.feature_flags = []
|
||||
success, msg = client.capture("distinct_id", "python test event")
|
||||
client.flush()
|
||||
self.assertTrue(success)
|
||||
self.assertFalse(self.failed)
|
||||
assert "$feature/beta-feature" not in msg["properties"]
|
||||
assert "$feature/beta-feature-local" not in msg["properties"]
|
||||
assert "$feature/false-flag" not in msg["properties"]
|
||||
assert "$active_feature_flags" not in msg["properties"]
|
||||
|
||||
@mock.patch("posthog.client.decide")
|
||||
def test_dont_override_capture_with_local_flags(self, patch_decide):
|
||||
patch_decide.return_value = {"featureFlags": {"beta-feature": "random-variant"}}
|
||||
client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, personal_api_key=FAKE_TEST_API_KEY)
|
||||
|
||||
multivariate_flag = {
|
||||
"id": 1,
|
||||
"name": "Beta Feature",
|
||||
"key": "beta-feature-local",
|
||||
"is_simple_flag": False,
|
||||
"active": True,
|
||||
"rollout_percentage": 100,
|
||||
"filters": {
|
||||
"groups": [
|
||||
{
|
||||
"properties": [
|
||||
{"key": "email", "type": "person", "value": "test@posthog.com", "operator": "exact"}
|
||||
],
|
||||
"rollout_percentage": 100,
|
||||
},
|
||||
{
|
||||
"rollout_percentage": 50,
|
||||
},
|
||||
],
|
||||
"multivariate": {
|
||||
"variants": [
|
||||
{"key": "first-variant", "name": "First Variant", "rollout_percentage": 50},
|
||||
{"key": "second-variant", "name": "Second Variant", "rollout_percentage": 25},
|
||||
{"key": "third-variant", "name": "Third Variant", "rollout_percentage": 25},
|
||||
]
|
||||
},
|
||||
"payloads": {"first-variant": "some-payload", "third-variant": {"a": "json"}},
|
||||
},
|
||||
}
|
||||
basic_flag = {
|
||||
"id": 1,
|
||||
"name": "Beta Feature",
|
||||
"key": "person-flag",
|
||||
"is_simple_flag": True,
|
||||
"active": True,
|
||||
"filters": {
|
||||
"groups": [
|
||||
{
|
||||
"properties": [
|
||||
{
|
||||
"key": "region",
|
||||
"operator": "exact",
|
||||
"value": ["USA"],
|
||||
"type": "person",
|
||||
}
|
||||
],
|
||||
"rollout_percentage": 100,
|
||||
}
|
||||
],
|
||||
"payloads": {"true": 300},
|
||||
},
|
||||
}
|
||||
client.feature_flags = [multivariate_flag, basic_flag]
|
||||
|
||||
success, msg = client.capture(
|
||||
"distinct_id", "python test event", {"$feature/beta-feature-local": "my-custom-variant"}
|
||||
)
|
||||
client.flush()
|
||||
self.assertTrue(success)
|
||||
self.assertFalse(self.failed)
|
||||
|
||||
self.assertEqual(msg["event"], "python test event")
|
||||
self.assertTrue(isinstance(msg["timestamp"], str))
|
||||
self.assertIsNone(msg.get("uuid"))
|
||||
self.assertEqual(msg["distinct_id"], "distinct_id")
|
||||
self.assertEqual(msg["properties"]["$lib"], "posthog-python")
|
||||
self.assertEqual(msg["properties"]["$lib_version"], VERSION)
|
||||
self.assertEqual(msg["properties"]["$feature/beta-feature-local"], "my-custom-variant")
|
||||
self.assertEqual(msg["properties"]["$active_feature_flags"], ["beta-feature-local"])
|
||||
assert "$feature/beta-feature" not in msg["properties"]
|
||||
assert "$feature/person-flag" not in msg["properties"]
|
||||
|
||||
self.assertEqual(patch_decide.call_count, 0)
|
||||
|
||||
@mock.patch("posthog.client.decide")
|
||||
def test_basic_capture_with_feature_flags_returns_active_only(self, patch_decide):
|
||||
@@ -131,6 +482,7 @@ class TestClient(unittest.TestCase):
|
||||
self.assertTrue(isinstance(msg["timestamp"], str))
|
||||
self.assertIsNone(msg.get("uuid"))
|
||||
self.assertEqual(msg["distinct_id"], "distinct_id")
|
||||
self.assertTrue(msg["properties"]["$geoip_disable"])
|
||||
self.assertEqual(msg["properties"]["$lib"], "posthog-python")
|
||||
self.assertEqual(msg["properties"]["$lib_version"], VERSION)
|
||||
self.assertEqual(msg["properties"]["$feature/beta-feature"], "random-variant")
|
||||
@@ -138,6 +490,58 @@ class TestClient(unittest.TestCase):
|
||||
self.assertEqual(msg["properties"]["$active_feature_flags"], ["beta-feature", "alpha-feature"])
|
||||
|
||||
self.assertEqual(patch_decide.call_count, 1)
|
||||
patch_decide.assert_called_with(
|
||||
"random_key",
|
||||
"https://us.i.posthog.com",
|
||||
timeout=3,
|
||||
distinct_id="distinct_id",
|
||||
groups={},
|
||||
person_properties=None,
|
||||
group_properties=None,
|
||||
disable_geoip=True,
|
||||
)
|
||||
|
||||
@mock.patch("posthog.client.decide")
|
||||
def test_basic_capture_with_feature_flags_and_disable_geoip_returns_correctly(self, patch_decide):
|
||||
patch_decide.return_value = {
|
||||
"featureFlags": {"beta-feature": "random-variant", "alpha-feature": True, "off-feature": False}
|
||||
}
|
||||
|
||||
client = Client(
|
||||
FAKE_TEST_API_KEY,
|
||||
host="https://app.posthog.com",
|
||||
on_error=self.set_fail,
|
||||
personal_api_key=FAKE_TEST_API_KEY,
|
||||
disable_geoip=True,
|
||||
feature_flags_request_timeout_seconds=12,
|
||||
)
|
||||
success, msg = client.capture("distinct_id", "python test event", send_feature_flags=True, disable_geoip=False)
|
||||
client.flush()
|
||||
self.assertTrue(success)
|
||||
self.assertFalse(self.failed)
|
||||
|
||||
self.assertEqual(msg["event"], "python test event")
|
||||
self.assertTrue(isinstance(msg["timestamp"], str))
|
||||
self.assertIsNone(msg.get("uuid"))
|
||||
self.assertTrue("$geoip_disable" not in msg["properties"])
|
||||
self.assertEqual(msg["distinct_id"], "distinct_id")
|
||||
self.assertEqual(msg["properties"]["$lib"], "posthog-python")
|
||||
self.assertEqual(msg["properties"]["$lib_version"], VERSION)
|
||||
self.assertEqual(msg["properties"]["$feature/beta-feature"], "random-variant")
|
||||
self.assertEqual(msg["properties"]["$feature/alpha-feature"], True)
|
||||
self.assertEqual(msg["properties"]["$active_feature_flags"], ["beta-feature", "alpha-feature"])
|
||||
|
||||
self.assertEqual(patch_decide.call_count, 1)
|
||||
patch_decide.assert_called_with(
|
||||
"random_key",
|
||||
"https://us.i.posthog.com",
|
||||
timeout=12,
|
||||
distinct_id="distinct_id",
|
||||
groups={},
|
||||
person_properties=None,
|
||||
group_properties=None,
|
||||
disable_geoip=False,
|
||||
)
|
||||
|
||||
@mock.patch("posthog.client.decide")
|
||||
def test_basic_capture_with_feature_flags_switched_off_doesnt_send_them(self, patch_decide):
|
||||
@@ -305,6 +709,26 @@ class TestClient(unittest.TestCase):
|
||||
"$group_set": {},
|
||||
"$lib": "posthog-python",
|
||||
"$lib_version": VERSION,
|
||||
"$geoip_disable": True,
|
||||
},
|
||||
)
|
||||
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))
|
||||
@@ -326,6 +750,36 @@ class TestClient(unittest.TestCase):
|
||||
"$group_set": {"trait": "value"},
|
||||
"$lib": "posthog-python",
|
||||
"$lib_version": VERSION,
|
||||
"$geoip_disable": True,
|
||||
},
|
||||
)
|
||||
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"},
|
||||
{"ip": "192.168.0.1"},
|
||||
datetime(2014, 9, 3),
|
||||
"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")
|
||||
@@ -477,6 +931,33 @@ class TestClient(unittest.TestCase):
|
||||
|
||||
self.assertEqual(msg, "disabled")
|
||||
|
||||
@mock.patch("posthog.client.decide")
|
||||
def test_disabled_with_feature_flags(self, patch_decide):
|
||||
client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, disabled=True)
|
||||
|
||||
response = client.get_feature_flag("beta-feature", "12345")
|
||||
self.assertIsNone(response)
|
||||
patch_decide.assert_not_called()
|
||||
|
||||
response = client.feature_enabled("beta-feature", "12345")
|
||||
self.assertIsNone(response)
|
||||
patch_decide.assert_not_called()
|
||||
|
||||
response = client.get_all_flags("12345")
|
||||
self.assertIsNone(response)
|
||||
patch_decide.assert_not_called()
|
||||
|
||||
response = client.get_feature_flag_payload("key", "12345")
|
||||
self.assertIsNone(response)
|
||||
patch_decide.assert_not_called()
|
||||
|
||||
response = client.get_all_flags_and_payloads("12345")
|
||||
self.assertEqual(response, {"featureFlags": None, "featureFlagPayloads": None})
|
||||
patch_decide.assert_not_called()
|
||||
|
||||
# no capture calls
|
||||
self.assertTrue(client.queue.empty())
|
||||
|
||||
def test_enabled_to_disabled(self):
|
||||
client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, disabled=False)
|
||||
success, msg = client.capture("distinct_id", "python test event")
|
||||
@@ -494,6 +975,74 @@ class TestClient(unittest.TestCase):
|
||||
|
||||
self.assertEqual(msg, "disabled")
|
||||
|
||||
def test_disable_geoip_default_on_events(self):
|
||||
client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, disable_geoip=True)
|
||||
_, capture_msg = client.capture("distinct_id", "python test event")
|
||||
client.flush()
|
||||
self.assertEqual(capture_msg["properties"]["$geoip_disable"], True)
|
||||
|
||||
_, identify_msg = client.identify("distinct_id", {"trait": "value"})
|
||||
client.flush()
|
||||
self.assertEqual(identify_msg["properties"]["$geoip_disable"], True)
|
||||
|
||||
def test_disable_geoip_override_on_events(self):
|
||||
client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, disable_geoip=False)
|
||||
_, capture_msg = client.set("distinct_id", {"a": "b", "c": "d"}, disable_geoip=True)
|
||||
client.flush()
|
||||
self.assertEqual(capture_msg["properties"]["$geoip_disable"], True)
|
||||
|
||||
_, identify_msg = client.page("distinct_id", "http://a.com", {"trait": "value"}, disable_geoip=False)
|
||||
client.flush()
|
||||
self.assertEqual("$geoip_disable" not in identify_msg["properties"], True)
|
||||
|
||||
def test_disable_geoip_method_overrides_init_on_events(self):
|
||||
client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, disable_geoip=True)
|
||||
_, msg = client.capture("distinct_id", "python test event", disable_geoip=False)
|
||||
client.flush()
|
||||
self.assertTrue("$geoip_disable" not in msg["properties"])
|
||||
|
||||
@mock.patch("posthog.client.decide")
|
||||
def test_disable_geoip_default_on_decide(self, patch_decide):
|
||||
patch_decide.return_value = {
|
||||
"featureFlags": {"beta-feature": "random-variant", "alpha-feature": True, "off-feature": False}
|
||||
}
|
||||
client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, disable_geoip=False)
|
||||
client.get_feature_flag("random_key", "some_id", disable_geoip=True)
|
||||
patch_decide.assert_called_with(
|
||||
"random_key",
|
||||
"https://us.i.posthog.com",
|
||||
timeout=3,
|
||||
distinct_id="some_id",
|
||||
groups={},
|
||||
person_properties={"distinct_id": "some_id"},
|
||||
group_properties={},
|
||||
disable_geoip=True,
|
||||
)
|
||||
patch_decide.reset_mock()
|
||||
client.feature_enabled("random_key", "feature_enabled_distinct_id", disable_geoip=True)
|
||||
patch_decide.assert_called_with(
|
||||
"random_key",
|
||||
"https://us.i.posthog.com",
|
||||
timeout=3,
|
||||
distinct_id="feature_enabled_distinct_id",
|
||||
groups={},
|
||||
person_properties={"distinct_id": "feature_enabled_distinct_id"},
|
||||
group_properties={},
|
||||
disable_geoip=True,
|
||||
)
|
||||
patch_decide.reset_mock()
|
||||
client.get_all_flags_and_payloads("all_flags_payloads_id")
|
||||
patch_decide.assert_called_with(
|
||||
"random_key",
|
||||
"https://us.i.posthog.com",
|
||||
timeout=3,
|
||||
distinct_id="all_flags_payloads_id",
|
||||
groups={},
|
||||
person_properties={"distinct_id": "all_flags_payloads_id"},
|
||||
group_properties={},
|
||||
disable_geoip=False,
|
||||
)
|
||||
|
||||
@mock.patch("posthog.client.Poller")
|
||||
@mock.patch("posthog.client.get")
|
||||
def test_call_identify_fails(self, patch_get, patch_poll):
|
||||
@@ -505,3 +1054,70 @@ class TestClient(unittest.TestCase):
|
||||
client.feature_flags = [{"key": "example", "is_simple_flag": False}]
|
||||
|
||||
self.assertFalse(client.feature_enabled("example", "distinct_id"))
|
||||
|
||||
@mock.patch("posthog.client.decide")
|
||||
def test_default_properties_get_added_properly(self, patch_decide):
|
||||
patch_decide.return_value = {
|
||||
"featureFlags": {"beta-feature": "random-variant", "alpha-feature": True, "off-feature": False}
|
||||
}
|
||||
client = Client(FAKE_TEST_API_KEY, host="http://app2.posthog.com", on_error=self.set_fail, disable_geoip=False)
|
||||
client.get_feature_flag(
|
||||
"random_key",
|
||||
"some_id",
|
||||
groups={"company": "id:5", "instance": "app.posthog.com"},
|
||||
person_properties={"x1": "y1"},
|
||||
group_properties={"company": {"x": "y"}},
|
||||
)
|
||||
patch_decide.assert_called_with(
|
||||
"random_key",
|
||||
"http://app2.posthog.com",
|
||||
timeout=3,
|
||||
distinct_id="some_id",
|
||||
groups={"company": "id:5", "instance": "app.posthog.com"},
|
||||
person_properties={"distinct_id": "some_id", "x1": "y1"},
|
||||
group_properties={
|
||||
"company": {"$group_key": "id:5", "x": "y"},
|
||||
"instance": {"$group_key": "app.posthog.com"},
|
||||
},
|
||||
disable_geoip=False,
|
||||
)
|
||||
|
||||
patch_decide.reset_mock()
|
||||
client.get_feature_flag(
|
||||
"random_key",
|
||||
"some_id",
|
||||
groups={"company": "id:5", "instance": "app.posthog.com"},
|
||||
person_properties={"distinct_id": "override"},
|
||||
group_properties={
|
||||
"company": {
|
||||
"$group_key": "group_override",
|
||||
}
|
||||
},
|
||||
)
|
||||
patch_decide.assert_called_with(
|
||||
"random_key",
|
||||
"http://app2.posthog.com",
|
||||
timeout=3,
|
||||
distinct_id="some_id",
|
||||
groups={"company": "id:5", "instance": "app.posthog.com"},
|
||||
person_properties={"distinct_id": "override"},
|
||||
group_properties={
|
||||
"company": {"$group_key": "group_override"},
|
||||
"instance": {"$group_key": "app.posthog.com"},
|
||||
},
|
||||
disable_geoip=False,
|
||||
)
|
||||
|
||||
patch_decide.reset_mock()
|
||||
# test nones
|
||||
client.get_all_flags_and_payloads("some_id", groups={}, person_properties=None, group_properties=None)
|
||||
patch_decide.assert_called_with(
|
||||
"random_key",
|
||||
"http://app2.posthog.com",
|
||||
timeout=3,
|
||||
distinct_id="some_id",
|
||||
groups={},
|
||||
person_properties={"distinct_id": "some_id"},
|
||||
group_properties={},
|
||||
disable_geoip=False,
|
||||
)
|
||||
|
||||
@@ -145,15 +145,20 @@ class TestConsumer(unittest.TestCase):
|
||||
def test_max_batch_size(self):
|
||||
q = Queue()
|
||||
consumer = Consumer(q, TEST_API_KEY, flush_at=100000, flush_interval=3)
|
||||
track = {"type": "track", "event": "python event", "distinct_id": "distinct_id"}
|
||||
properties = {}
|
||||
for n in range(0, 500):
|
||||
properties[str(n)] = "one_long_property_value_to_build_a_big_event"
|
||||
track = {"type": "track", "event": "python event", "distinct_id": "distinct_id", "properties": properties}
|
||||
msg_size = len(json.dumps(track).encode())
|
||||
# number of messages in a maximum-size batch
|
||||
n_msgs = int(475000 / msg_size)
|
||||
# Let's capture 8MB of data to trigger two batches
|
||||
n_msgs = int(8_000_000 / msg_size)
|
||||
|
||||
def mock_post_fn(_, data, **kwargs):
|
||||
res = mock.Mock()
|
||||
res.status_code = 200
|
||||
self.assertTrue(len(data.encode()) < 500000, "batch size (%d) exceeds 500KB limit" % len(data.encode()))
|
||||
request_size = len(data.encode())
|
||||
# Batches close after the first message bringing it bigger than BATCH_SIZE_LIMIT, let's add 10% of margin
|
||||
self.assertTrue(request_size < (5 * 1024 * 1024) * 1.1, "batch size (%d) higher than limit" % request_size)
|
||||
return res
|
||||
|
||||
with mock.patch("posthog.request._session.post", side_effect=mock_post_fn) as mock_post:
|
||||
|
||||
@@ -0,0 +1,63 @@
|
||||
import subprocess
|
||||
import sys
|
||||
from textwrap import dedent
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def test_excepthook(tmpdir):
|
||||
app = tmpdir.join("app.py")
|
||||
app.write(
|
||||
dedent(
|
||||
"""
|
||||
from posthog import Posthog
|
||||
posthog = Posthog('phc_x', host='https://eu.i.posthog.com', enable_exception_autocapture=True, debug=True, on_error=lambda e, batch: print('error handling batch: ', e, batch))
|
||||
|
||||
# frame_value = "LOL"
|
||||
|
||||
1/0
|
||||
"""
|
||||
)
|
||||
)
|
||||
|
||||
with pytest.raises(subprocess.CalledProcessError) as excinfo:
|
||||
subprocess.check_output([sys.executable, str(app)], stderr=subprocess.STDOUT)
|
||||
|
||||
output = excinfo.value.output
|
||||
|
||||
assert b"ZeroDivisionError" in output
|
||||
assert b"LOL" in output
|
||||
assert b"DEBUG:posthog:data uploaded successfully" in output
|
||||
assert (
|
||||
b'"$exception_list": [{"mechanism": {"type": "generic", "handled": true}, "module": null, "type": "ZeroDivisionError", "value": "division by zero", "stacktrace": {"frames": [{"platform": "python", "filename": "app.py", "abs_path"'
|
||||
in output
|
||||
)
|
||||
|
||||
|
||||
def test_trying_to_use_django_integration(tmpdir):
|
||||
app = tmpdir.join("app.py")
|
||||
app.write(
|
||||
dedent(
|
||||
"""
|
||||
from posthog import Posthog, Integrations
|
||||
posthog = Posthog('phc_x', host='https://eu.i.posthog.com', enable_exception_autocapture=True, exception_autocapture_integrations=[Integrations.Django], debug=True, on_error=lambda e, batch: print('error handling batch: ', e, batch))
|
||||
|
||||
# frame_value = "LOL"
|
||||
|
||||
1/0
|
||||
"""
|
||||
)
|
||||
)
|
||||
|
||||
with pytest.raises(subprocess.CalledProcessError) as excinfo:
|
||||
subprocess.check_output([sys.executable, str(app)], stderr=subprocess.STDOUT)
|
||||
|
||||
output = excinfo.value.output
|
||||
|
||||
assert b"ZeroDivisionError" in output
|
||||
assert b"LOL" in output
|
||||
assert b"DEBUG:posthog:data uploaded successfully" in output
|
||||
assert (
|
||||
b'"$exception_list": [{"mechanism": {"type": "generic", "handled": true}, "module": null, "type": "ZeroDivisionError", "value": "division by zero", "stacktrace": {"frames": [{"platform": "python", "filename": "app.py", "abs_path"'
|
||||
in output
|
||||
)
|
||||
@@ -6,7 +6,7 @@ from dateutil import parser, tz
|
||||
from freezegun import freeze_time
|
||||
|
||||
from posthog.client import Client
|
||||
from posthog.feature_flags import InconclusiveMatchError, match_property
|
||||
from posthog.feature_flags import InconclusiveMatchError, match_property, relative_date_parse_for_feature_flag_matching
|
||||
from posthog.request import APIError
|
||||
from posthog.test.test_utils import FAKE_TEST_API_KEY
|
||||
|
||||
@@ -960,6 +960,60 @@ class TestLocalEvaluation(unittest.TestCase):
|
||||
self.assertEqual(patch_decide.call_count, 0)
|
||||
self.assertEqual(patch_capture.call_count, 0)
|
||||
|
||||
@mock.patch("posthog.client.decide")
|
||||
@mock.patch("posthog.client.get")
|
||||
def test_feature_flags_local_evaluation_None_values(self, patch_get, patch_decide):
|
||||
client = Client(FAKE_TEST_API_KEY, personal_api_key=FAKE_TEST_API_KEY)
|
||||
client.feature_flags = [
|
||||
{
|
||||
id: 1,
|
||||
"name": "Beta Feature",
|
||||
"key": "beta-feature",
|
||||
"is_simple_flag": True,
|
||||
"active": True,
|
||||
"filters": {
|
||||
"groups": [
|
||||
{
|
||||
"variant": None,
|
||||
"properties": [
|
||||
{"key": "latestBuildVersion", "type": "person", "value": ".+", "operator": "regex"},
|
||||
{"key": "latestBuildVersionMajor", "type": "person", "value": "23", "operator": "gt"},
|
||||
{"key": "latestBuildVersionMinor", "type": "person", "value": "31", "operator": "gt"},
|
||||
{"key": "latestBuildVersionPatch", "type": "person", "value": "0", "operator": "gt"},
|
||||
],
|
||||
"rollout_percentage": 100,
|
||||
}
|
||||
],
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
feature_flag_match = client.get_feature_flag(
|
||||
"beta-feature",
|
||||
"some-distinct-id",
|
||||
person_properties={
|
||||
"latestBuildVersion": None,
|
||||
"latestBuildVersionMajor": None,
|
||||
"latestBuildVersionMinor": None,
|
||||
"latestBuildVersionPatch": None,
|
||||
},
|
||||
)
|
||||
|
||||
self.assertEqual(feature_flag_match, False)
|
||||
self.assertEqual(patch_decide.call_count, 0)
|
||||
self.assertEqual(patch_get.call_count, 0)
|
||||
|
||||
feature_flag_match = client.get_feature_flag(
|
||||
"beta-feature",
|
||||
"some-distinct-id",
|
||||
person_properties={
|
||||
"latestBuildVersion": "24.32..1",
|
||||
"latestBuildVersionMajor": "24",
|
||||
"latestBuildVersionMinor": "32",
|
||||
"latestBuildVersionPatch": "1",
|
||||
},
|
||||
)
|
||||
|
||||
@mock.patch("posthog.client.decide")
|
||||
@mock.patch("posthog.client.get")
|
||||
def test_feature_flags_local_evaluation_for_cohorts(self, patch_get, patch_decide):
|
||||
@@ -1578,9 +1632,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"}
|
||||
@@ -1595,6 +1650,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):
|
||||
@@ -1714,7 +1771,7 @@ class TestMatchProperties(unittest.TestCase):
|
||||
self.assertTrue(match_property(property_a, {"key": "value"}))
|
||||
self.assertTrue(match_property(property_a, {"key": "value2"}))
|
||||
self.assertTrue(match_property(property_a, {"key": ""}))
|
||||
self.assertTrue(match_property(property_a, {"key": None}))
|
||||
self.assertFalse(match_property(property_a, {"key": None}))
|
||||
|
||||
with self.assertRaises(InconclusiveMatchError):
|
||||
match_property(property_a, {"key2": "value"})
|
||||
@@ -1775,7 +1832,8 @@ class TestMatchProperties(unittest.TestCase):
|
||||
|
||||
self.assertFalse(match_property(property_a, {"key": 0}))
|
||||
self.assertFalse(match_property(property_a, {"key": -1}))
|
||||
self.assertFalse(match_property(property_a, {"key": "23"}))
|
||||
# now we handle type mismatches so this should be true
|
||||
self.assertTrue(match_property(property_a, {"key": "23"}))
|
||||
|
||||
property_b = self.property(key="key", value=1, operator="lt")
|
||||
self.assertTrue(match_property(property_b, {"key": 0}))
|
||||
@@ -1792,7 +1850,8 @@ class TestMatchProperties(unittest.TestCase):
|
||||
|
||||
self.assertFalse(match_property(property_c, {"key": 0}))
|
||||
self.assertFalse(match_property(property_c, {"key": -1}))
|
||||
self.assertFalse(match_property(property_c, {"key": "3"}))
|
||||
# now we handle type mismatches so this should be true
|
||||
self.assertTrue(match_property(property_c, {"key": "3"}))
|
||||
|
||||
property_d = self.property(key="key", value="43", operator="lte")
|
||||
self.assertTrue(match_property(property_d, {"key": "41"}))
|
||||
@@ -1801,6 +1860,21 @@ class TestMatchProperties(unittest.TestCase):
|
||||
|
||||
self.assertFalse(match_property(property_d, {"key": "44"}))
|
||||
self.assertFalse(match_property(property_d, {"key": 44}))
|
||||
self.assertTrue(match_property(property_d, {"key": 42}))
|
||||
|
||||
property_e = self.property(key="key", value="30", operator="lt")
|
||||
self.assertTrue(match_property(property_e, {"key": "29"}))
|
||||
|
||||
# depending on the type of override, we adjust type comparison
|
||||
self.assertTrue(match_property(property_e, {"key": "100"}))
|
||||
self.assertFalse(match_property(property_e, {"key": 100}))
|
||||
|
||||
property_f = self.property(key="key", value="123aloha", operator="gt")
|
||||
self.assertFalse(match_property(property_f, {"key": "123"}))
|
||||
self.assertFalse(match_property(property_f, {"key": 122}))
|
||||
|
||||
# this turns into a string comparison
|
||||
self.assertTrue(match_property(property_f, {"key": 129}))
|
||||
|
||||
def test_match_property_date_operators(self):
|
||||
property_a = self.property(key="key", value="2022-05-01", operator="is_date_before")
|
||||
@@ -1854,6 +1928,303 @@ class TestMatchProperties(unittest.TestCase):
|
||||
self.assertTrue(match_property(property_d, {"key": "2022-04-05 11:34:11 +00:00"}))
|
||||
self.assertFalse(match_property(property_d, {"key": "2022-04-05 11:34:13 +00:00"}))
|
||||
|
||||
@freeze_time("2022-05-01")
|
||||
def test_match_property_relative_date_operators(self):
|
||||
property_a = self.property(key="key", value="-6h", operator="is_date_before")
|
||||
self.assertTrue(match_property(property_a, {"key": "2022-03-01"}))
|
||||
self.assertTrue(match_property(property_a, {"key": "2022-04-30"}))
|
||||
self.assertTrue(match_property(property_a, {"key": datetime.datetime(2022, 4, 30, 1, 2, 3)}))
|
||||
# false because date comparison, instead of datetime, so reduces to same date
|
||||
self.assertFalse(match_property(property_a, {"key": datetime.date(2022, 4, 30)}))
|
||||
|
||||
self.assertFalse(match_property(property_a, {"key": datetime.datetime(2022, 4, 30, 19, 2, 3)}))
|
||||
self.assertTrue(
|
||||
match_property(
|
||||
property_a,
|
||||
{"key": datetime.datetime(2022, 4, 30, 1, 2, 3, tzinfo=tz.gettz("Europe/Madrid"))},
|
||||
)
|
||||
)
|
||||
self.assertTrue(match_property(property_a, {"key": parser.parse("2022-04-30")}))
|
||||
self.assertFalse(match_property(property_a, {"key": "2022-05-30"}))
|
||||
|
||||
# Can't be a number
|
||||
with self.assertRaises(InconclusiveMatchError):
|
||||
match_property(property_a, {"key": 1})
|
||||
|
||||
# can't be invalid string
|
||||
with self.assertRaises(InconclusiveMatchError):
|
||||
match_property(property_a, {"key": "abcdef"})
|
||||
|
||||
property_b = self.property(key="key", value="1h", operator="is_date_after")
|
||||
self.assertTrue(match_property(property_b, {"key": "2022-05-02"}))
|
||||
self.assertTrue(match_property(property_b, {"key": "2022-05-30"}))
|
||||
self.assertTrue(match_property(property_b, {"key": datetime.datetime(2022, 5, 30)}))
|
||||
self.assertTrue(match_property(property_b, {"key": parser.parse("2022-05-30")}))
|
||||
self.assertFalse(match_property(property_b, {"key": "2022-04-30"}))
|
||||
|
||||
# can't be invalid string
|
||||
with self.assertRaises(InconclusiveMatchError):
|
||||
self.assertFalse(match_property(property_b, {"key": "abcdef"}))
|
||||
|
||||
# Invalid flag property
|
||||
property_c = self.property(key="key", value=1234, operator="is_date_after")
|
||||
|
||||
with self.assertRaises(InconclusiveMatchError):
|
||||
self.assertFalse(match_property(property_c, {"key": 1}))
|
||||
|
||||
# parsed as 1234-05-01 for some reason?
|
||||
self.assertTrue(match_property(property_c, {"key": "2022-05-30"}))
|
||||
|
||||
# # Timezone aware property
|
||||
property_d = self.property(key="key", value="12d", operator="is_date_before")
|
||||
self.assertFalse(match_property(property_d, {"key": "2022-05-30"}))
|
||||
|
||||
self.assertTrue(match_property(property_d, {"key": "2022-03-30"}))
|
||||
self.assertTrue(match_property(property_d, {"key": "2022-04-05 12:34:11+01:00"}))
|
||||
self.assertTrue(match_property(property_d, {"key": "2022-04-19 01:34:11+02:00"}))
|
||||
|
||||
self.assertFalse(match_property(property_d, {"key": "2022-04-19 02:00:01+02:00"}))
|
||||
|
||||
# Try all possible relative dates
|
||||
property_e = self.property(key="key", value="1h", operator="is_date_before")
|
||||
self.assertFalse(match_property(property_e, {"key": "2022-05-01 00:00:00"}))
|
||||
self.assertTrue(match_property(property_e, {"key": "2022-04-30 22:00:00"}))
|
||||
|
||||
property_f = self.property(key="key", value="-1d", operator="is_date_before")
|
||||
self.assertTrue(match_property(property_f, {"key": "2022-04-29 23:59:00"}))
|
||||
self.assertFalse(match_property(property_f, {"key": "2022-04-30 00:00:01"}))
|
||||
|
||||
property_g = self.property(key="key", value="1w", operator="is_date_before")
|
||||
self.assertTrue(match_property(property_g, {"key": "2022-04-23 00:00:00"}))
|
||||
self.assertFalse(match_property(property_g, {"key": "2022-04-24 00:00:00"}))
|
||||
self.assertFalse(match_property(property_g, {"key": "2022-04-24 00:00:01"}))
|
||||
|
||||
property_h = self.property(key="key", value="1m", operator="is_date_before")
|
||||
self.assertTrue(match_property(property_h, {"key": "2022-03-01 00:00:00"}))
|
||||
self.assertFalse(match_property(property_h, {"key": "2022-04-05 00:00:00"}))
|
||||
|
||||
property_i = self.property(key="key", value="1y", operator="is_date_before")
|
||||
self.assertTrue(match_property(property_i, {"key": "2021-04-28 00:00:00"}))
|
||||
self.assertFalse(match_property(property_i, {"key": "2021-05-01 00:00:01"}))
|
||||
|
||||
property_j = self.property(key="key", value="122h", operator="is_date_after")
|
||||
self.assertTrue(match_property(property_j, {"key": "2022-05-01 00:00:00"}))
|
||||
self.assertFalse(match_property(property_j, {"key": "2022-04-23 01:00:00"}))
|
||||
|
||||
property_k = self.property(key="key", value="2d", operator="is_date_after")
|
||||
self.assertTrue(match_property(property_k, {"key": "2022-05-01 00:00:00"}))
|
||||
self.assertTrue(match_property(property_k, {"key": "2022-04-29 00:00:01"}))
|
||||
self.assertFalse(match_property(property_k, {"key": "2022-04-29 00:00:00"}))
|
||||
|
||||
property_l = self.property(key="key", value="-02w", operator="is_date_after")
|
||||
self.assertTrue(match_property(property_l, {"key": "2022-05-01 00:00:00"}))
|
||||
self.assertFalse(match_property(property_l, {"key": "2022-04-16 00:00:00"}))
|
||||
|
||||
property_m = self.property(key="key", value="1m", operator="is_date_after")
|
||||
self.assertTrue(match_property(property_m, {"key": "2022-04-01 00:00:01"}))
|
||||
self.assertFalse(match_property(property_m, {"key": "2022-04-01 00:00:00"}))
|
||||
|
||||
property_n = self.property(key="key", value="1y", operator="is_date_after")
|
||||
self.assertTrue(match_property(property_n, {"key": "2022-05-01 00:00:00"}))
|
||||
self.assertTrue(match_property(property_n, {"key": "2021-05-01 00:00:01"}))
|
||||
self.assertFalse(match_property(property_n, {"key": "2021-05-01 00:00:00"}))
|
||||
self.assertFalse(match_property(property_n, {"key": "2021-04-30 00:00:00"}))
|
||||
self.assertFalse(match_property(property_n, {"key": "2021-03-01 12:13:00"}))
|
||||
|
||||
def test_none_property_value_with_all_operators(self):
|
||||
property_a = self.property(key="key", value="none", operator="is_not")
|
||||
self.assertFalse(match_property(property_a, {"key": None}))
|
||||
self.assertTrue(match_property(property_a, {"key": "non"}))
|
||||
|
||||
property_b = self.property(key="key", value=None, operator="is_set")
|
||||
self.assertFalse(match_property(property_b, {"key": None}))
|
||||
|
||||
property_c = self.property(key="key", value="no", operator="icontains")
|
||||
self.assertFalse(match_property(property_c, {"key": None}))
|
||||
self.assertFalse(match_property(property_c, {"key": "smh"}))
|
||||
|
||||
property_d = self.property(key="key", value="No", operator="regex")
|
||||
self.assertFalse(match_property(property_d, {"key": None}))
|
||||
|
||||
property_d_lower_case = self.property(key="key", value="no", operator="regex")
|
||||
self.assertFalse(match_property(property_d_lower_case, {"key": None}))
|
||||
|
||||
property_e = self.property(key="key", value=1, operator="gt")
|
||||
self.assertFalse(match_property(property_e, {"key": None}))
|
||||
|
||||
property_f = self.property(key="key", value=1, operator="lt")
|
||||
self.assertFalse(match_property(property_f, {"key": None}))
|
||||
|
||||
property_g = self.property(key="key", value="xyz", operator="gte")
|
||||
self.assertFalse(match_property(property_g, {"key": None}))
|
||||
|
||||
property_h = self.property(key="key", value="Oo", operator="lte")
|
||||
self.assertFalse(match_property(property_h, {"key": None}))
|
||||
|
||||
property_i = self.property(key="key", value="2022-05-01", operator="is_date_before")
|
||||
self.assertFalse(match_property(property_i, {"key": None}))
|
||||
|
||||
property_j = self.property(key="key", value="2022-05-01", operator="is_date_after")
|
||||
self.assertFalse(match_property(property_j, {"key": None}))
|
||||
|
||||
property_k = self.property(key="key", value="2022-05-01", operator="is_date_before")
|
||||
with self.assertRaises(InconclusiveMatchError):
|
||||
self.assertFalse(match_property(property_k, {"key": "random"}))
|
||||
|
||||
def test_unknown_operator(self):
|
||||
property_a = self.property(key="key", value="2022-05-01", operator="is_unknown")
|
||||
with self.assertRaises(InconclusiveMatchError) as exception_context:
|
||||
match_property(property_a, {"key": "random"})
|
||||
self.assertEqual(str(exception_context.exception), "Unknown operator is_unknown")
|
||||
|
||||
|
||||
class TestRelativeDateParsing(unittest.TestCase):
|
||||
def test_invalid_input(self):
|
||||
with freeze_time("2020-01-01T12:01:20.1340Z"):
|
||||
assert relative_date_parse_for_feature_flag_matching("1") is None
|
||||
assert relative_date_parse_for_feature_flag_matching("1x") is None
|
||||
assert relative_date_parse_for_feature_flag_matching("1.2y") is None
|
||||
assert relative_date_parse_for_feature_flag_matching("1z") is None
|
||||
assert relative_date_parse_for_feature_flag_matching("1s") is None
|
||||
assert relative_date_parse_for_feature_flag_matching("123344000.134m") is None
|
||||
assert relative_date_parse_for_feature_flag_matching("bazinga") is None
|
||||
assert relative_date_parse_for_feature_flag_matching("000bello") is None
|
||||
assert relative_date_parse_for_feature_flag_matching("000hello") is None
|
||||
|
||||
assert relative_date_parse_for_feature_flag_matching("000h") is not None
|
||||
assert relative_date_parse_for_feature_flag_matching("1000h") is not None
|
||||
|
||||
def test_overflow(self):
|
||||
assert relative_date_parse_for_feature_flag_matching("1000000h") is None
|
||||
assert relative_date_parse_for_feature_flag_matching("100000000000000000y") is None
|
||||
|
||||
def test_hour_parsing(self):
|
||||
with freeze_time("2020-01-01T12:01:20.1340Z"):
|
||||
assert relative_date_parse_for_feature_flag_matching("1h") == datetime.datetime(
|
||||
2020, 1, 1, 11, 1, 20, 134000, tzinfo=tz.gettz("UTC")
|
||||
)
|
||||
assert relative_date_parse_for_feature_flag_matching("2h") == datetime.datetime(
|
||||
2020, 1, 1, 10, 1, 20, 134000, tzinfo=tz.gettz("UTC")
|
||||
)
|
||||
assert relative_date_parse_for_feature_flag_matching("24h") == datetime.datetime(
|
||||
2019, 12, 31, 12, 1, 20, 134000, tzinfo=tz.gettz("UTC")
|
||||
)
|
||||
assert relative_date_parse_for_feature_flag_matching("30h") == datetime.datetime(
|
||||
2019, 12, 31, 6, 1, 20, 134000, tzinfo=tz.gettz("UTC")
|
||||
)
|
||||
assert relative_date_parse_for_feature_flag_matching("48h") == datetime.datetime(
|
||||
2019, 12, 30, 12, 1, 20, 134000, tzinfo=tz.gettz("UTC")
|
||||
)
|
||||
|
||||
assert relative_date_parse_for_feature_flag_matching(
|
||||
"24h"
|
||||
) == relative_date_parse_for_feature_flag_matching("1d")
|
||||
assert relative_date_parse_for_feature_flag_matching(
|
||||
"48h"
|
||||
) == relative_date_parse_for_feature_flag_matching("2d")
|
||||
|
||||
def test_day_parsing(self):
|
||||
with freeze_time("2020-01-01T12:01:20.1340Z"):
|
||||
assert relative_date_parse_for_feature_flag_matching("1d") == datetime.datetime(
|
||||
2019, 12, 31, 12, 1, 20, 134000, tzinfo=tz.gettz("UTC")
|
||||
)
|
||||
assert relative_date_parse_for_feature_flag_matching("2d") == datetime.datetime(
|
||||
2019, 12, 30, 12, 1, 20, 134000, tzinfo=tz.gettz("UTC")
|
||||
)
|
||||
assert relative_date_parse_for_feature_flag_matching("7d") == datetime.datetime(
|
||||
2019, 12, 25, 12, 1, 20, 134000, tzinfo=tz.gettz("UTC")
|
||||
)
|
||||
assert relative_date_parse_for_feature_flag_matching("14d") == datetime.datetime(
|
||||
2019, 12, 18, 12, 1, 20, 134000, tzinfo=tz.gettz("UTC")
|
||||
)
|
||||
assert relative_date_parse_for_feature_flag_matching("30d") == datetime.datetime(
|
||||
2019, 12, 2, 12, 1, 20, 134000, tzinfo=tz.gettz("UTC")
|
||||
)
|
||||
|
||||
assert relative_date_parse_for_feature_flag_matching("7d") == relative_date_parse_for_feature_flag_matching(
|
||||
"1w"
|
||||
)
|
||||
|
||||
def test_week_parsing(self):
|
||||
with freeze_time("2020-01-01T12:01:20.1340Z"):
|
||||
assert relative_date_parse_for_feature_flag_matching("1w") == datetime.datetime(
|
||||
2019, 12, 25, 12, 1, 20, 134000, tzinfo=tz.gettz("UTC")
|
||||
)
|
||||
assert relative_date_parse_for_feature_flag_matching("2w") == datetime.datetime(
|
||||
2019, 12, 18, 12, 1, 20, 134000, tzinfo=tz.gettz("UTC")
|
||||
)
|
||||
assert relative_date_parse_for_feature_flag_matching("4w") == datetime.datetime(
|
||||
2019, 12, 4, 12, 1, 20, 134000, tzinfo=tz.gettz("UTC")
|
||||
)
|
||||
assert relative_date_parse_for_feature_flag_matching("8w") == datetime.datetime(
|
||||
2019, 11, 6, 12, 1, 20, 134000, tzinfo=tz.gettz("UTC")
|
||||
)
|
||||
|
||||
assert relative_date_parse_for_feature_flag_matching("1m") == datetime.datetime(
|
||||
2019, 12, 1, 12, 1, 20, 134000, tzinfo=tz.gettz("UTC")
|
||||
)
|
||||
assert relative_date_parse_for_feature_flag_matching("4w") != relative_date_parse_for_feature_flag_matching(
|
||||
"1m"
|
||||
)
|
||||
|
||||
def test_month_parsing(self):
|
||||
with freeze_time("2020-01-01T12:01:20.1340Z"):
|
||||
assert relative_date_parse_for_feature_flag_matching("1m") == datetime.datetime(
|
||||
2019, 12, 1, 12, 1, 20, 134000, tzinfo=tz.gettz("UTC")
|
||||
)
|
||||
assert relative_date_parse_for_feature_flag_matching("2m") == datetime.datetime(
|
||||
2019, 11, 1, 12, 1, 20, 134000, tzinfo=tz.gettz("UTC")
|
||||
)
|
||||
assert relative_date_parse_for_feature_flag_matching("4m") == datetime.datetime(
|
||||
2019, 9, 1, 12, 1, 20, 134000, tzinfo=tz.gettz("UTC")
|
||||
)
|
||||
assert relative_date_parse_for_feature_flag_matching("8m") == datetime.datetime(
|
||||
2019, 5, 1, 12, 1, 20, 134000, tzinfo=tz.gettz("UTC")
|
||||
)
|
||||
|
||||
assert relative_date_parse_for_feature_flag_matching("1y") == datetime.datetime(
|
||||
2019, 1, 1, 12, 1, 20, 134000, tzinfo=tz.gettz("UTC")
|
||||
)
|
||||
assert relative_date_parse_for_feature_flag_matching(
|
||||
"12m"
|
||||
) == relative_date_parse_for_feature_flag_matching("1y")
|
||||
|
||||
with freeze_time("2020-04-03T00:00:00"):
|
||||
assert relative_date_parse_for_feature_flag_matching("1m") == datetime.datetime(
|
||||
2020, 3, 3, 0, 0, 0, tzinfo=tz.gettz("UTC")
|
||||
)
|
||||
assert relative_date_parse_for_feature_flag_matching("2m") == datetime.datetime(
|
||||
2020, 2, 3, 0, 0, 0, tzinfo=tz.gettz("UTC")
|
||||
)
|
||||
assert relative_date_parse_for_feature_flag_matching("4m") == datetime.datetime(
|
||||
2019, 12, 3, 0, 0, 0, tzinfo=tz.gettz("UTC")
|
||||
)
|
||||
assert relative_date_parse_for_feature_flag_matching("8m") == datetime.datetime(
|
||||
2019, 8, 3, 0, 0, 0, tzinfo=tz.gettz("UTC")
|
||||
)
|
||||
|
||||
assert relative_date_parse_for_feature_flag_matching("1y") == datetime.datetime(
|
||||
2019, 4, 3, 0, 0, 0, tzinfo=tz.gettz("UTC")
|
||||
)
|
||||
assert relative_date_parse_for_feature_flag_matching(
|
||||
"12m"
|
||||
) == relative_date_parse_for_feature_flag_matching("1y")
|
||||
|
||||
def test_year_parsing(self):
|
||||
with freeze_time("2020-01-01T12:01:20.1340Z"):
|
||||
assert relative_date_parse_for_feature_flag_matching("1y") == datetime.datetime(
|
||||
2019, 1, 1, 12, 1, 20, 134000, tzinfo=tz.gettz("UTC")
|
||||
)
|
||||
assert relative_date_parse_for_feature_flag_matching("2y") == datetime.datetime(
|
||||
2018, 1, 1, 12, 1, 20, 134000, tzinfo=tz.gettz("UTC")
|
||||
)
|
||||
assert relative_date_parse_for_feature_flag_matching("4y") == datetime.datetime(
|
||||
2016, 1, 1, 12, 1, 20, 134000, tzinfo=tz.gettz("UTC")
|
||||
)
|
||||
assert relative_date_parse_for_feature_flag_matching("8y") == datetime.datetime(
|
||||
2012, 1, 1, 12, 1, 20, 134000, tzinfo=tz.gettz("UTC")
|
||||
)
|
||||
|
||||
|
||||
class TestCaptureCalls(unittest.TestCase):
|
||||
@mock.patch.object(Client, "capture")
|
||||
@@ -1888,8 +2259,14 @@ class TestCaptureCalls(unittest.TestCase):
|
||||
patch_capture.assert_called_with(
|
||||
"some-distinct-id",
|
||||
"$feature_flag_called",
|
||||
{"$feature_flag": "complex-flag", "$feature_flag_response": True, "locally_evaluated": True},
|
||||
{
|
||||
"$feature_flag": "complex-flag",
|
||||
"$feature_flag_response": True,
|
||||
"locally_evaluated": True,
|
||||
"$feature/complex-flag": True,
|
||||
},
|
||||
groups={},
|
||||
disable_geoip=None,
|
||||
)
|
||||
patch_capture.reset_mock()
|
||||
|
||||
@@ -1912,8 +2289,14 @@ class TestCaptureCalls(unittest.TestCase):
|
||||
patch_capture.assert_called_with(
|
||||
"some-distinct-id2",
|
||||
"$feature_flag_called",
|
||||
{"$feature_flag": "complex-flag", "$feature_flag_response": True, "locally_evaluated": True},
|
||||
{
|
||||
"$feature_flag": "complex-flag",
|
||||
"$feature_flag_response": True,
|
||||
"locally_evaluated": True,
|
||||
"$feature/complex-flag": True,
|
||||
},
|
||||
groups={},
|
||||
disable_geoip=None,
|
||||
)
|
||||
patch_capture.reset_mock()
|
||||
|
||||
@@ -1944,8 +2327,139 @@ class TestCaptureCalls(unittest.TestCase):
|
||||
patch_capture.assert_called_with(
|
||||
"some-distinct-id2",
|
||||
"$feature_flag_called",
|
||||
{"$feature_flag": "decide-flag", "$feature_flag_response": "decide-value", "locally_evaluated": False},
|
||||
{
|
||||
"$feature_flag": "decide-flag",
|
||||
"$feature_flag_response": "decide-value",
|
||||
"locally_evaluated": False,
|
||||
"$feature/decide-flag": "decide-value",
|
||||
},
|
||||
groups={"organization": "org1"},
|
||||
disable_geoip=None,
|
||||
)
|
||||
|
||||
@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",
|
||||
"is_simple_flag": False,
|
||||
"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):
|
||||
patch_decide.return_value = {"featureFlags": {"decide-flag": "decide-value"}}
|
||||
client = Client(FAKE_TEST_API_KEY, personal_api_key=FAKE_TEST_API_KEY, disable_geoip=True)
|
||||
client.feature_flags = [
|
||||
{
|
||||
"id": 1,
|
||||
"name": "Beta Feature",
|
||||
"key": "complex-flag",
|
||||
"is_simple_flag": False,
|
||||
"active": True,
|
||||
"filters": {
|
||||
"groups": [
|
||||
{
|
||||
"properties": [{"key": "region", "value": "USA"}],
|
||||
"rollout_percentage": 100,
|
||||
}
|
||||
],
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
client.get_feature_flag(
|
||||
"complex-flag",
|
||||
"some-distinct-id",
|
||||
person_properties={"region": "USA", "name": "Aloha"},
|
||||
disable_geoip=False,
|
||||
)
|
||||
|
||||
patch_capture.assert_called_with(
|
||||
"some-distinct-id",
|
||||
"$feature_flag_called",
|
||||
{
|
||||
"$feature_flag": "complex-flag",
|
||||
"$feature_flag_response": True,
|
||||
"locally_evaluated": True,
|
||||
"$feature/complex-flag": True,
|
||||
},
|
||||
groups={},
|
||||
disable_geoip=False,
|
||||
)
|
||||
|
||||
@mock.patch("posthog.client.MAX_DICT_SIZE", 100)
|
||||
@@ -1977,8 +2491,14 @@ class TestCaptureCalls(unittest.TestCase):
|
||||
patch_capture.assert_called_with(
|
||||
distinct_id,
|
||||
"$feature_flag_called",
|
||||
{"$feature_flag": "complex-flag", "$feature_flag_response": True, "locally_evaluated": True},
|
||||
{
|
||||
"$feature_flag": "complex-flag",
|
||||
"$feature_flag_response": True,
|
||||
"locally_evaluated": True,
|
||||
"$feature/complex-flag": True,
|
||||
},
|
||||
groups={},
|
||||
disable_geoip=None,
|
||||
)
|
||||
|
||||
self.assertEqual(len(client.distinct_ids_feature_flags_reported), i % 100 + 1)
|
||||
|
||||
@@ -6,6 +6,10 @@ from posthog import Posthog
|
||||
class TestModule(unittest.TestCase):
|
||||
posthog = None
|
||||
|
||||
def _assert_enqueue_result(self, result):
|
||||
self.assertEqual(type(result[0]), bool)
|
||||
self.assertEqual(type(result[1]), dict)
|
||||
|
||||
def failed(self):
|
||||
self.failed = True
|
||||
|
||||
@@ -22,15 +26,18 @@ class TestModule(unittest.TestCase):
|
||||
self.assertRaises(Exception, self.posthog.capture)
|
||||
|
||||
def test_track(self):
|
||||
self.posthog.capture("distinct_id", "python module event")
|
||||
res = self.posthog.capture("distinct_id", "python module event")
|
||||
self._assert_enqueue_result(res)
|
||||
self.posthog.flush()
|
||||
|
||||
def test_identify(self):
|
||||
self.posthog.identify("distinct_id", {"email": "user@email.com"})
|
||||
res = self.posthog.identify("distinct_id", {"email": "user@email.com"})
|
||||
self._assert_enqueue_result(res)
|
||||
self.posthog.flush()
|
||||
|
||||
def test_alias(self):
|
||||
self.posthog.alias("previousId", "distinct_id")
|
||||
res = self.posthog.alias("previousId", "distinct_id")
|
||||
self._assert_enqueue_result(res)
|
||||
self.posthog.flush()
|
||||
|
||||
def test_page(self):
|
||||
|
||||
@@ -2,9 +2,10 @@ import json
|
||||
import unittest
|
||||
from datetime import date, datetime
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from posthog.request import DatetimeSerializer, batch_post
|
||||
from posthog.request import DatetimeSerializer, batch_post, determine_server_host
|
||||
from posthog.test.test_utils import TEST_API_KEY
|
||||
|
||||
|
||||
@@ -42,3 +43,26 @@ class TestRequests(unittest.TestCase):
|
||||
batch_post(
|
||||
"key", batch=[{"distinct_id": "distinct_id", "event": "python event", "type": "track"}], timeout=0.0001
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"host, expected",
|
||||
[
|
||||
("https://t.posthog.com", "https://t.posthog.com"),
|
||||
("https://t.posthog.com/", "https://t.posthog.com/"),
|
||||
("t.posthog.com", "t.posthog.com"),
|
||||
("t.posthog.com/", "t.posthog.com/"),
|
||||
("https://us.posthog.com.rg.proxy.com", "https://us.posthog.com.rg.proxy.com"),
|
||||
("app.posthog.com", "app.posthog.com"),
|
||||
("eu.posthog.com", "eu.posthog.com"),
|
||||
("https://app.posthog.com", "https://us.i.posthog.com"),
|
||||
("https://eu.posthog.com", "https://eu.i.posthog.com"),
|
||||
("https://us.posthog.com", "https://us.i.posthog.com"),
|
||||
("https://app.posthog.com/", "https://us.i.posthog.com"),
|
||||
("https://eu.posthog.com/", "https://eu.i.posthog.com"),
|
||||
("https://us.posthog.com/", "https://us.i.posthog.com"),
|
||||
(None, "https://us.i.posthog.com"),
|
||||
],
|
||||
)
|
||||
def test_routing_to_custom_host(host, expected):
|
||||
assert determine_server_host(host) == expected
|
||||
|
||||
+1
-1
@@ -1,4 +1,4 @@
|
||||
VERSION = "2.5.0"
|
||||
VERSION = "3.9.0"
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(VERSION, end="") # noqa: T201
|
||||
|
||||
@@ -13,6 +13,7 @@ Including another URLconf
|
||||
1. Import the include() function: from django.urls import include, path
|
||||
2. Add a URL to urlpatterns: path('blog/', include('blog.urls'))
|
||||
"""
|
||||
|
||||
from django.contrib import admin
|
||||
from django.urls import path
|
||||
|
||||
|
||||
@@ -1,2 +1,5 @@
|
||||
[bdist_wheel]
|
||||
universal = 1
|
||||
|
||||
[tool:pytest]
|
||||
asyncio_mode = auto
|
||||
|
||||
@@ -14,7 +14,13 @@ 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",
|
||||
]
|
||||
|
||||
extras_require = {
|
||||
"dev": [
|
||||
@@ -24,8 +30,25 @@ extras_require = {
|
||||
"flake8-print",
|
||||
"pre-commit",
|
||||
],
|
||||
"test": ["mock>=2.0.0", "freezegun==0.3.15", "pylint", "flake8", "coverage", "pytest"],
|
||||
"test": [
|
||||
"mock>=2.0.0",
|
||||
"freezegun==0.3.15",
|
||||
"pylint",
|
||||
"flake8",
|
||||
"coverage",
|
||||
"pytest",
|
||||
"pytest-timeout",
|
||||
"pytest-asyncio",
|
||||
"django",
|
||||
"openai",
|
||||
"anthropic",
|
||||
"langgraph",
|
||||
"langchain-community>=0.2.0",
|
||||
"langchain-openai>=0.2.0",
|
||||
"langchain-anthropic>=0.2.0",
|
||||
],
|
||||
"sentry": ["sentry-sdk", "django"],
|
||||
"langchain": ["langchain>=0.2.0"],
|
||||
}
|
||||
|
||||
setup(
|
||||
@@ -37,7 +60,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,
|
||||
@@ -60,5 +92,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",
|
||||
],
|
||||
)
|
||||
|
||||
+13
-1
@@ -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",
|
||||
],
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user