Compare commits

..
424 Commits
Author SHA1 Message Date
zeekay 81fb40cde3 corona→corona: academic Corona now only in lp-220-p3q-corona 2026-06-11 10:25:54 -07:00
Hanzo AI 001ac2f65b chore: symlink AGENTS.md and CLAUDE.md to LLM.md 2026-04-01 14:13:28 -07:00
Hanzo AI b5f5dade80 chore: rename docker-compose.yml to compose.yml 2026-03-03 14:00:00 -08:00
Hanzo Dev 29a74b090a rebrand: remove all compat aliases, zero posthog refs
- Remove Posthog = Insights alias from __init__.py
- Remove PosthogContextMiddleware alias from django.py
- Remove PostHogTracingProcessor alias from processor.py
- Change $lib from "posthog-python" to "insights-python"
- Change ingestion URLs from posthog.com to insights.hanzo.ai
- Rename all posthog_* kwargs to insights_* across AI wrappers
- Rename __posthog_exception_captured to __insights_exception_captured
- Rename posthog_context_stack contextvar to insights_context_stack
- Rename posthog🎏 Redis prefix to insights🎏
- Rename $$_posthog_redacted_* sentinels to $$_insights_redacted_*
- Remove POSTHOG_MW_* Django settings fallback, X-POSTHOG-* headers
- Rename Prompts(posthog=) param to Prompts(client=)
- Update APP_ENDPOINT to us.insights.hanzo.ai
- Update all tests, examples, docs, mypy config
2026-03-13 20:29:13 -07:00
Hanzo Dev 24318fe57c rebrand: Posthog->Insights, package hanzo-insights
- Rename main class Posthog -> Insights (Posthog kept as alias)
- Rename PosthogContextMiddleware -> InsightsContextMiddleware (alias kept)
- Rename PostHogTracingProcessor -> InsightsTracingProcessor (alias kept)
- Add `insights/` re-export package so `from insights import Insights` works
- Update all imports from `posthog` to `hanzo_insights` across source and tests
- Update docstrings, comments, error messages, user agent string
- Update README, example.py, .env.example, Makefile, LLM.md
- Django middleware now supports INSIGHTS_MW_* settings (POSTHOG_MW_* still works)
- Django middleware accepts X-INSIGHTS-* headers (X-POSTHOG-* still works)
- Keep protocol-level values ($lib, ingestion URLs, sentinel strings) for server compat
- Keep posthog_* parameter names in AI wrappers for API compat
- All 681 tests pass
2026-03-13 19:58:09 -07:00
Hanzo Dev e14fe3824d docs: add LLM.md project guide 2026-03-11 10:32:50 -07:00
Hanzo Dev 894c008b5d chore: rename package from hanzoanalytics to hanzo-insights
Package name: hanzo-insights (import as hanzo_insights)
2026-03-06 22:44:19 -08:00
Hanzo Dev 867d7d0aed Rename package from posthog/posthoganalytics to hanzoanalytics
Full rebrand: module directory, pyproject.toml, setup.py, all imports.
2026-03-06 22:42:05 -08:00
Radu RaiceaandGitHub afb978e1f8 feat(llma): support prompt versions in prompts sdk (#454)
* feat(llma): support prompt versions in prompts sdk

* fix(llma): enforce clear_cache version requires name
2026-03-06 10:29:07 +01:00
github-actions[bot] d73abc6d93 chore: Release v7.9.7 2026-03-05 22:09:29 +00:00
a96ed6616c Add warning log for local flag evaluation cold start (#452)
* Add warning log when local flag evaluation called before flags loaded

When feature_enabled() is called with only_evaluate_locally=True before
flag definitions are fetched, the SDK silently returns None. This adds a
warning log so users can diagnose the issue immediately.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* Move cold start warning to only fire for only_evaluate_locally=True

The warning was in _locally_evaluate_flag which runs for all flag
evaluations, including those that fall back to server-side evaluation.
Move it to the caller where only_evaluate_locally is known, so it only
fires when the caller explicitly opted out of the server fallback.

* Narrow cold start warning to only fire when flags were never fetched

Use `is None` instead of `not` to avoid firing when flags are loaded
but empty (401, 402, no personal_api_key), which already have their
own specific error logs.

* add changeset

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-05 13:36:05 -08:00
fb8677bd18 fix(llma): use distinct_id from outer context if not provided (#449)
* fix(llma): use distinct_id from outer context if not provided

* fix(llma): distinct_id from context is now explicitly passed to capture method

* fix(llma): fix $process_person_profile with outer context distinct_id, add tests

- Fix personless check to consider outer context distinct_id (not just the
  explicit param), so events from users who set distinct_id via outer context
  are not incorrectly marked as personless.
- Fix typo: "district_id" -> "distinct_id" in comments.
- Add test coverage for distinct_id resolution: no id (personless), explicit
  param, outer context, and explicit overriding outer context.

* chore: add sampo changeset for distinct_id context fix

* style: ruff format

---------

Co-authored-by: Andrew Maguire <andrewm4894@gmail.com>
2026-03-05 15:11:40 +00:00
github-actions[bot] 369b530db9 chore: Release v7.9.6 2026-03-02 21:28:45 +00:00
451cbc85be refactor: add PROPERTY_OPERATORS constant for match_property (#448)
* feat: add semver targeting support to local flag evaluation

Implement 9 semver comparison operators (semver_eq, semver_neq, semver_gt, semver_gte, semver_lt, semver_lte, semver_tilde, semver_caret, semver_wildcard) for feature flag local evaluation. Uses regex-based parsing that matches the server-side sortableSemver behavior to handle v-prefix, whitespace, pre-release suffixes, and non-standard version formats.

Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com>

* fix: guard against ReDoS in semver regex parsing

Add input length limit before regex search to prevent polynomial
backtracking on adversarial input (CodeQL py/polynomial-redos).

* fix: replace regex with string parsing to resolve ReDoS warning

Replace SEMVER_EXTRACT_RE regex with simple string splitting to
eliminate nested quantifiers that CodeQL flagged as polynomial-redos.

* refactor: inline semver operator tuple to match existing patterns

* refactor: add PROPERTY_OPERATORS constant for match_property

Extract all operator strings into a single source-of-truth tuple and
validate against it early in match_property, replacing the fallthrough
at the end of the function.

* refactor: split PROPERTY_OPERATORS into composable sub-groups

Break the flat tuple into category-specific tuples (EQUALITY_OPERATORS,
STRING_OPERATORS, etc.) that compose into PROPERTY_OPERATORS via
concatenation. The semver dispatch code now references
SEMVER_OPERATORS and SEMVER_COMPARISON_OPERATORS instead of
repeating the full operator lists inline.

* fix: add unreachable fallthrough to satisfy mypy return check

* add release

---------

Co-authored-by: Claude Haiku 4.5 <noreply@anthropic.com>
2026-03-02 21:12:09 +00:00
github-actions[bot] af050f03f8 chore: Release v7.9.5 2026-03-02 20:53:53 +00:00
37beabd1d3 feat: add semver targeting support to local flag evaluation (#447)
* feat: add semver targeting support to local flag evaluation

Implement 9 semver comparison operators (semver_eq, semver_neq, semver_gt, semver_gte, semver_lt, semver_lte, semver_tilde, semver_caret, semver_wildcard) for feature flag local evaluation. Uses regex-based parsing that matches the server-side sortableSemver behavior to handle v-prefix, whitespace, pre-release suffixes, and non-standard version formats.

Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com>

* fix: guard against ReDoS in semver regex parsing

Add input length limit before regex search to prevent polynomial
backtracking on adversarial input (CodeQL py/polynomial-redos).

* fix: replace regex with string parsing to resolve ReDoS warning

Replace SEMVER_EXTRACT_RE regex with simple string splitting to
eliminate nested quantifiers that CodeQL flagged as polynomial-redos.

* refactor: inline semver operator tuple to match existing patterns

* add changeset

---------

Co-authored-by: Claude Haiku 4.5 <noreply@anthropic.com>
2026-03-02 20:52:37 +00:00
github-actions[bot] 9214babf55 chore: Release v7.9.4 2026-02-25 15:28:29 +00:00
Carlos MarchalandGitHub a2ea0bbfe7 fix: revert manual release and add sampo changeset for ai_tokens_source (#445) 2026-02-25 16:25:53 +01:00
Andrew MaguireandGitHub 91f8c39c53 feat(llma): add $ai_tokens_source property to detect token value overrides (#444)
* feat: add $ai_tokens_source property to detect token value overrides

When users pass token properties (e.g. $ai_input_tokens) via
posthog_properties, these override the SDK-computed values. This new
$ai_tokens_source property ("sdk" or "passthrough") lets us distinguish
whether token values came from the SDK or were externally injected,
which is critical for diagnosing cost calculation discrepancies.

* chore: bump version to 7.9.4

* chore: add changelog entry for 7.9.4

* chore: fix ruff formatting

* chore: remove unused pytest import
2026-02-25 13:38:53 +00:00
Michael BiancoandGitHub c689530d82 fix: use $ip not $ip_addess (#356) 2026-02-20 07:46:13 +01:00
github-actions[bot] ad2ef1ef39 chore: Release v7.9.3 2026-02-18 22:20:09 +00:00
Rafael AudibertandGitHub 54a0d98d6f Small fixes for python publishing (#441)
* fix: Avoid setting dynamic version

Version is now fixed because of sampo, so we can get rid of this

* feat: add changeset

* docs: Add new RELEASING section to README
2026-02-18 22:17:26 +00:00
github-actions[bot] 6ba2d8af81 chore: Release v7.9.2 2026-02-18 22:05:00 +00:00
Rafael AudibertandGitHub d6c97b371f chore: Migrate releases to sampo (#398)
* chore: Migrate releases to `sampo`

This is much closer to what we have in `posthog-js`, let's see if it's a good thing!

There's still a lot to do before deploying this:
- updating CI
- updating README with instructions

* Add sampo changeset

* chore: Update  to relase Python via Slack + sampo

* Update release.yml

* fix: Use pyproject.toml version as source of truth
2026-02-18 19:02:19 -03:00
Radu RaiceaandGitHub ef7c513b59 fix(ai): bind prompt reads to project token (#433)
* fix(ai): bind prompt reads to project token

* chore(release): bump version to 7.8.7
2026-02-17 16:58:59 +00:00
c877700b43 chore: release v7.9.0 (#434)
chore: bump version to 7.9.0

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-17 10:32:05 +01:00
AndersandGitHub 0fa03d328e feat: Support device_id as bucketing identifier for local evaluation (#424)
* feat: Support device_id as bucketing identifier for local evaluation

Add support for `bucketing_identifier` field on feature flags to allow
using `device_id` instead of `distinct_id` for hashing/bucketing in
local evaluation.
2026-02-16 11:05:09 +01:00
AndersandGitHub 0639f0f31a test(flags): make wrong-key load_feature_flags deterministic (#432) 2026-02-13 15:53:22 +01:00
Aleksander BłaszkiewiczandGitHub f8a5f109c1 feat: limit max number of items in collection to scan (#430)
* feat: limit max number of items in collection to scan

* feat: changelog

* fix: format

* feat: test

* feat: replace entire collection instead of truncating
2026-02-11 14:59:11 +01:00
Aleksander BłaszkiewiczandGitHub 005709ff76 feat: further optimize code variables regex search (#429)
* feat: initial

* fix: ruff
2026-02-09 23:59:03 +01:00
Aleksander BłaszkiewiczandGitHub 100e0c08a4 fix: long variables pattern matching (#428) 2026-02-09 17:45:23 +01:00
AndersandGitHub 2791ee0175 fix: Retry on 408 and respect Retry-After header (#426)
* fix: Retry on 408 and respect Retry-After header

408 (Request Timeout) was incorrectly treated as a non-retryable client
error. Retry-After response headers were ignored during backoff. Replace
backoff library usage with a manual retry loop that honours Retry-After
when present and falls back to exponential backoff otherwise.

* fix: Parse HTTP-date Retry-After values

Retry-After can be seconds or an HTTP-date per RFC 7231. Fall back to
email.utils.parsedate_to_datetime when the numeric parse fails.

* fix: Don't retry on unclassifiable APIError status

When APIError.status is "N/A" (no HTTP status), treat it as
non-retryable to avoid unexpected retry loops on errors the SDK
cannot classify.

* test: Add retry delay tests for Retry-After and exponential backoff

Verify time.sleep is called with the Retry-After value when present,
uses exponential backoff (2^attempt) when absent, and that 408 is
retried.
2026-02-09 12:15:49 +00:00
Aleksander BłaszkiewiczandGitHub 07a5261f0f fix: openai image sanitization (#425) 2026-02-06 14:15:57 +01:00
f409515bfa fix: openAI input image sanitization (#384)
Co-authored-by: Aleksander Błaszkiewicz <kqmdjc8@gmail.com>
2026-02-06 13:53:02 +01:00
Gabriel GrinbergandGitHub c50b7b2e78 Fix feature flag 401 errors causing HTTP request storm (#422)
* Fix feature flag 401 errors causing HTTP request storm

Set feature_flags = [] on 401 error to prevent repeated requests.

* Clear flag_cache, group_type_mapping, cohorts on 401
2026-02-04 10:31:10 -05:00
Radu RaiceaandGitHub cbc74eb1d4 fix(llma): default prompts url (#423) 2026-02-04 15:10:00 +00:00
Radu RaiceaandGitHub 786ef4d612 fix(llma): small fixes for prompt management (#420)
* fix(llma): small fixes for prompt management

* fix(llma): tests

* fix(llma): tests
2026-02-04 09:49:19 +02:00
José SequeiraandGitHub cd351653c0 feat: SDK Compliance (#397)
* feat: SDK Compliance
2026-01-30 16:12:43 +01:00
Radu RaiceaandGitHub ea320eab5a feat(llma): add prompt management (#417)
* feat(llma): add prompt management

* chore(llma): bump version

* fix(llma): use SDK session with retry logic for prompt fetching

Use _get_session() from posthog/request.py instead of raw requests.get()
to benefit from the SDK's existing retry configuration on transient
network failures.
2026-01-30 08:43:04 -05:00
ea56c52be4 feat(llma): pass raw provider usage metadata for backend cost calculations (#411)
* feat: pass raw provider usage metadata for backend cost calculations

Add raw_usage field to TokenUsage type to capture raw provider usage metadata (OpenAI, Anthropic, Gemini). This enables the backend to extract modality-specific token counts (text vs image vs audio) for accurate cost calculations.

- Add raw_usage field to TokenUsage TypedDict
- Update all provider converters to capture raw usage:
  - OpenAI: capture response.usage and chunk usage
  - Anthropic: capture usage from message_start and message_delta events
  - Gemini: capture usage_metadata from responses and chunks
- Pass raw usage as $ai_usage property in PostHog events
- Update merge_usage_stats to handle raw_usage in both modes
- Add tests verifying $ai_usage is captured for all providers

Backend will extract provider-specific details and delete $ai_usage after processing to avoid bloating properties.

Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

* fix: add serialize_raw_usage helper to ensure JSON serializability

Address PR review feedback from @andrewm4894:

1. **Serialization**: Add serialize_raw_usage() helper with fallback chain:
   - .model_dump() for Pydantic models (OpenAI/Anthropic)
   - .to_dict() for protobuf-like objects
   - vars() for simple objects
   - str() as last resort
   This ensures we never pass unserializable objects to PostHog client.

2. **Data loss prevention**: Change from replacing to merging raw_usage in
   incremental mode. For Anthropic streaming, message_start has input token
   details and message_delta has output token details - merging preserves
   both instead of losing input data.

3. **Test coverage**: Enhanced tests to verify:
   - JSON serializability with json.dumps()
   - Expected structure of raw_usage dicts
   - Coverage for both non-streaming and streaming modes
   - Fixed Gemini test mocks to return proper dicts from model_dump()

Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

* refactor: move raw_usage serialization from utils to converters

Address PR feedback from @andrewm4894 - serialize in converters, not utils.

**Problem:**
Utils was receiving raw Pydantic/protobuf objects and serializing them,
which meant provider-specific knowledge leaked into generic code.

**Solution:**
Move serialization into converters where provider context exists:

Converters (NEW):
- OpenAI: serialize_raw_usage(response.usage) → dict
- Anthropic: serialize_raw_usage(event.usage) → dict
- Gemini: serialize_raw_usage(metadata) → dict

Utils (SIMPLIFIED):
- Just passes dicts through, no serialization needed
- Merge operations work with dicts only

**Benefits:**
1. Type correctness: raw_usage is always Dict[str, Any]
2. Separation of concerns: converters handle provider formats
3. Fail fast: serialization errors in converters with context
4. Cleaner abstraction: utils doesn't know about Pydantic/protobuf

**Flow:**
Provider object → Converter serializes → dict → Utils → PostHog

Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

* fix: add type annotation for current_raw to satisfy mypy

Fix mypy error: "Need type annotation for 'current_raw'"

Extract value first, then apply explicit type annotation with ternary
conditional to satisfy mypy's type checker.

Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com>
2026-01-28 11:51:02 +02:00
Andrew MaguireandGitHub 6d3e769256 feat(ai): add OpenAI Agents SDK integration (#408)
* feat(ai): add OpenAI Agents SDK integration

Add PostHogTracingProcessor that implements the OpenAI Agents SDK
TracingProcessor interface to capture agent traces in PostHog.

- Maps GenerationSpanData to $ai_generation events
- Maps FunctionSpanData, AgentSpanData, HandoffSpanData, GuardrailSpanData
  to $ai_span events with appropriate types
- Supports privacy mode, groups, and custom properties
- Includes instrument() helper for one-liner setup
- 22 unit tests covering all span types

* feat(openai-agents): add $ai_group_id support for linking conversation traces

- Capture group_id from trace and include as $ai_group_id on all events
- Add _get_group_id() helper to retrieve group_id from trace metadata
- Pass group_id through all span handlers (generation, function, agent, handoff, guardrail, response, custom, audio, mcp, generic)
- Enables linking multiple traces in the same conversation thread

* feat(openai-agents): add enhanced span properties

- Add $ai_total_tokens to generation and response spans (required by PostHog cost reporting)
- Add $ai_error_type for cross-provider error categorization (model_behavior_error, user_error, input_guardrail_triggered, output_guardrail_triggered, max_turns_exceeded)
- Add $ai_output_choices to response spans for output content capture
- Add audio pass-through properties for voice spans:
  - first_content_at (time to first audio byte)
  - audio_input_format / audio_output_format
  - model_config
  - $ai_input for TTS text input
- Add comprehensive tests for all new properties

* Add $ai_framework property and standardize $ai_provider for OpenAI Agents

- Add $ai_framework="openai-agents" to all events for framework identification
- Standardize $ai_provider="openai" on all events (previously some used "openai_agents")
- Follows pattern from posthog-js where $ai_provider is the underlying LLM provider

* chore: bump version to 7.7.0 for OpenAI Agents SDK integration

* fix: add openai_agents package to setuptools config

Without this, the module is not included in the distribution
and users get an ImportError after pip install.

* fix: correct indentation in on_trace_start properties dict

* fix: prevent unbounded growth of span/trace tracking dicts

Add max entry limit and eviction for _span_start_times and
_trace_metadata dicts. If on_span_end or on_trace_end is never
called (e.g., due to an SDK exception), these dicts could grow
indefinitely in long-running processes.

* fix: resolve distinct_id from trace metadata in on_span_end

Previously on_span_end always called _get_distinct_id(None), which
meant callable distinct_id resolvers never received the trace object
for spans. Now the resolved distinct_id is stored at trace start and
looked up by trace_id during span end.

* refactor: extract _base_properties helper to reduce duplication

All span handlers repeated the same 6 base fields (trace_id, span_id,
parent_id, provider, framework, latency) plus the group_id conditional.
Extract into a shared helper to reduce ~100 lines of boilerplate.

* test: add missing edge case tests for openai agents processor

- test_generation_span_with_no_usage: zero tokens when usage is None
- test_generation_span_with_partial_usage: only input_tokens present
- test_error_type_categorization_by_type_field_only: type field without
  matching message content
- test_distinct_id_resolved_from_trace_for_spans: callable resolver
  uses trace context for span events
- test_eviction_of_stale_entries: memory leak prevention works

* fix: handle non-dict error_info in span error parsing

If span.error is a string instead of a dict, calling .get() would
raise AttributeError. Now falls back to str() for non-dict errors.

* style: apply ruff formatting

* style: replace lambda assignments with def (ruff E731)

* fix: restore full CHANGELOG.md history

The rebase conflict resolution accidentally truncated the changelog
to only the most recent entries. Restored all historical entries.

* fix: preserve personless mode for trace-id fallback distinct IDs

When no distinct_id is provided, _get_distinct_id falls back to
trace_id or "unknown". Since these are non-None strings, the
$process_person_profile=False check in _capture_event never fired,
creating unwanted person profiles keyed by trace IDs.

Track whether the user explicitly provided a distinct_id and use
that flag to control personless mode, matching the pattern used
by the langchain and openai integrations.

* fix: restore changelog history and fix personless mode edge cases

Two fixes from bot review:

1. CHANGELOG.md was accidentally truncated to 38 lines during rebase
   conflict resolution. Restored all 767 lines of history.

2. Personless mode now follows the same pattern as langchain/openai
   integrations: _get_distinct_id returns None when no user-provided
   ID is available, and callers set $process_person_profile=False
   before falling back to trace_id. This covers the edge case where
   a callable distinct_id returns None.

* fix: handle None token counts in generation span

Guard against input_tokens or output_tokens being None when computing
$ai_total_tokens to avoid TypeError.

* fix: check error_type_raw for all error categories

Check both error_type_raw and error_message for guardrail and
max_turns errors, consistent with how ModelBehaviorError and
UserError are already checked.

* fix: add type hints to instrument() function

* refactor: rename _safe_json to _ensure_serializable for clarity

The function validates JSON serializability and falls back to str(),
not serializes. Rename and update docstring to make the contract clear.

* refactor: emit $ai_trace at trace end instead of start

Move the $ai_trace event from on_trace_start to on_trace_end to
capture full metadata including latency, matching the LangChain
integration approach. on_trace_start now only stores metadata for
use by spans.

* style: fix ruff formatting

* fix: add TYPE_CHECKING imports for type hints in instrument()
2026-01-27 21:15:16 +00:00
AndersandGitHub ef6ef12d81 feat: add device_id to flags request payload (#407)
* feat: add device_id to flags request payload

Add device_id parameter to all feature flag methods, similar to how
distinct_id is handled. The device_id is included in the flags request
payload sent to the server.

- Add device_id parameter to Client methods and module-level functions
- Add context support via set_context_device_id() for automatic fallback
- Add tests for explicit device_id and context-based device_id
- Bump version to 7.6.0
2026-01-19 17:22:34 +01:00
Paul D'AmbraandGitHub 0dab1e7f34 Fix link formatting in CHANGELOG.md (#406)
Updated link formatting for clarity in changelog.
2026-01-08 19:59:24 -03:00
Paul D'AmbraandGitHub 982af288d2 chore: check for syntax warnings (#404) 2026-01-08 21:26:22 +00:00
cbdb9adf40 fix: avoid return from finally block to fix Python 3.14 SyntaxWarning (#361)
* fix: Avoid return from finally: block

This fixes a SyntaxWarning on Python 3.14.

```
❯ uvx --no-cache --python 3.14.0 --with posthog==6.7.11 python -c "import posthog"
Installed 11 packages in 5ms
.../lib/python3.14/site-packages/posthog/consumer.py:92: SyntaxWarning: 'return' in a 'finally' block
  return success
````

* add versioning info

---------

Co-authored-by: Paul D'Ambra <paul.dambra@gmail.com>
2026-01-08 21:18:11 +00:00
Tom PiccirelloandGitHub 664696f3eb feat: add CodeQL Advanced workflow (#405)
This is required to be able to run CodeQL against external contributors' PRs.
2026-01-07 22:17:54 -03:00
Paul D'AmbraGitHubClaudegreptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
ad27f5ef47 chore: add SDK version support table to README (#402)
* docs: add Python version support table to README

Add a table documenting which SDK versions introduced or dropped
support for different Python versions, based on CHANGELOG.md entries.

* fix: correct Python 3.14 support version to 7.4.3

* Apply suggestions from code review

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

---------

Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
2026-01-07 19:32:43 -03:00
Paul D'AmbraandGitHub c5d0280e3b chore: add a test to describe upload behaviour when there are errors (#403)
* chore: add a test to describe upload behaviour when there are errors

* refactor the test file

* add typehints to the test file
2026-01-07 22:16:15 +00:00
af034835aa feat: llma / error tracking integration (#376)
* feat: llma / error tracking integration

* capture all metadata in llm event

* instrument with contexts

* bump version

* indentation

* linting

* tests

* raise

* test: add exception capture integration tests for langchain

Add 6 tests covering the new LLMA + error tracking integration:
- capture_exception called on span/generation errors
- $exception_event_id added to AI events
- No capture when autocapture disabled
- AI properties passed to exception event
- Handles None return from capture_exception

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* fix: pass context tags to capture() for test compatibility

- Export get_tags() from posthog module
- Explicitly pass context tags to capture() in AI utils
- Fix $ai_model fallback to extract from response.model
- Fix ruff formatting in langchain test_callbacks.py

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* fix: disable auto-capture exceptions in LLM context

The new_context() defaults to capture_exceptions=True which would
auto-capture any exception regardless of enable_exception_autocapture
setting. This was inconsistent with LangChain callbacks which
explicitly check the setting.

Pass capture_exceptions=False to let exception handling be controlled
explicitly by the enable_exception_autocapture setting.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* fix: isolate LLM context with fresh=True to avoid tag inheritance

Use fresh=True to start with a clean context for each LLM call.
This avoids inheriting $ai_* tags from parent contexts which could
cause mismatched AI metadata due to the tag merge order bug in
contexts.py (parent tags incorrectly override child tags).

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* fix: correct tag merge order so child tags take precedence

The collect_tags() method had a bug where parent tags would overwrite
child tags, despite the comment saying the opposite. This fix ensures
child context tags properly override parent tags.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* refactor: remove fresh=True now that tag merge order is fixed

With the collect_tags() bug fixed, child tags properly override parent
tags. LLM events can now inherit useful parent context tags (request_id,
user info, etc.) while still having their $ai_* tags take precedence.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* test: add test for child tags overriding parent tags

Verifies that in non-fresh contexts, child tags properly override
parent tags with the same key while still inheriting other parent tags.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* chore: add TODO for OpenAI/Anthropic/Gemini exception capture

Document that exception capture needs to be added for the direct SDK
wrappers, similar to how it's implemented in LangChain callbacks.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

---------

Co-authored-by: David Newell <david@Mac.communityfibre.co.uk>
Co-authored-by: Andrew Maguire <andrewm4894@gmail.com>
Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
2026-01-07 13:11:20 +00:00
AndersandGitHub 649e04335d fix(release): restore GitHub release creation with gh CLI (#401)
The "Create GitHub release" step was broken in PR #386, which removed
the GITHUB_TOKEN env var from the actions/create-release action. The
action requires the token to be passed explicitly, so releases were
being published to PyPI but GitHub tags/releases were not created.

This replaces the archived actions/create-release@v1 with the gh CLI,
which is already used elsewhere in this workflow. The gh CLI properly
uses GH_TOKEN for authentication.
2026-01-06 14:27:28 +01:00
Carlos MarchalandGitHub f21a5984e3 fix: double counting anthropic langchain (#399) 2026-01-05 16:56:17 +01:00
Hugues PouillotandGitHub b64c2d97f9 feat: add in_app configuration for python SDK (#396)
* add in_app configuration for python SDK

* bump version

* add in_app_modules to init script as well
2025-12-22 12:02:48 +01:00
Andrew MaguireGitHubClaude Opus 4.5greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
16cd501674 fix(llma): extract model from response for OpenAI stored prompts (#395)
* fix: extract model from response for OpenAI stored prompts

When using OpenAI stored prompts, the model is defined in the OpenAI
dashboard rather than passed in the API request. This change adds a
fallback to extract the model from the response object when not
provided in kwargs.

Fixes PostHog/posthog#42861

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* Apply suggestion from @greptile-apps[bot]

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

* Apply suggestion from @greptile-apps[bot]

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

* test: add tests for model extraction fallback and bump to 7.4.1

- Add 8 tests covering model extraction from response for stored prompts
- Fix utils.py to add 'unknown' fallback for consistency
- Bump version to 7.4.1
- Update CHANGELOG.md

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* style: format utils.py with ruff

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* fix: remove 'unknown' fallback from non-streaming to match original behavior

Non-streaming originally returned None when model wasn't in kwargs.
Streaming keeps "unknown" fallback as that was the original behavior.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* test: add test for None model fallback in non-streaming

Verifies that non-streaming returns None (not "unknown") when model
is not available in kwargs or response, matching original behavior.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
2025-12-21 21:17:32 +00:00
Dustin ByrneandGitHub ddfa1ee93c fix: Respect send_feature_flags setting and deprecate send_feature_flag_events in get_feature_flag_payload (#391)
* fix: Respect `send_feature_flags` when local eval is enabled

* fix: Deprecate `send_feature_flag_events` in `get_feature_flag_payload`
2025-12-17 13:04:10 -05:00
Phil HaackandGitHub 4bbccef790 feat(flags): Add retry support for feature flag requests (#392)
* Add urllib3-based retry for feature flag requests

Use urllib3's built-in Retry mechanism for feature flag POST requests
instead of application-level retry logic. This is simpler and leverages
well-tested library code.

Key changes:
- Add `RETRY_STATUS_FORCELIST` = [408, 500, 502, 503, 504]
- Add `_build_flags_session()` with POST retries and `status_forcelist`
- Update `flags()` to use dedicated flags session
- Add tests for retry configuration and session usage

The flags session retries on:
- Network failures (connect/read errors)
- Transient server errors (408, 500, 502, 503, 504)

It does NOT retry on:
- 429 (rate limit) - need to wait, not hammer
- 402 (quota limit) - won't resolve with retries

* Make examples run without requiring personal api key

* Add integration tests for network retry behavior

Add tests that verify actual retry behavior, not just configuration:

- test_retries_on_503_then_succeeds: Spins up a local HTTP server that
  returns 503 twice then 200, verifying 3 requests are made
- test_connection_errors_are_retried: Verifies connection errors trigger
  retries by measuring elapsed time with backoff

Both tests use dynamically allocated ports for CI safety.

* Bump version to 7.4.0
2025-12-16 23:41:38 +00:00
Phil HaackandGitHub 6fe1102993 feature: Add $feature_flag_error property to track flag evaluation failures (#390)
* Add $feature_flag_error property to track flag evaluation failures

Track errors in feature flag evaluation by adding a `$feature_flag_error` property to the `$feature_flag_called` event.

* Refactor requests exception imports through request.py

Export RequestsTimeout and RequestsConnectionError from posthog/request.py
to keep all requests library imports in one place and avoid mypy issues.

* Address PR review feedback

- Fix fallback logic to only trigger on actual exceptions, not when
  errors_while_computing or flag_missing is set from a successful API response
- Change log.exception() to log.warning() for expected operational errors
  (quota limits, timeouts, connection errors, API errors) to reduce log noise
- Keep log.exception() only for truly unexpected errors (unknown_error)
- Extract stale cache fallback into _get_stale_flag_fallback() helper method

* Add tests for stale cache fallback and error absence

- Add TestFeatureFlagErrorWithStaleCacheFallback test class with 4 tests:
  - test_timeout_error_returns_stale_cached_value
  - test_connection_error_returns_stale_cached_value
  - test_api_error_returns_stale_cached_value
  - test_error_without_cache_returns_none

- Add negative assertions to verify $feature_flag_error is absent on success:
  - test_get_feature_flag_result_boolean_local_evaluation
  - test_get_feature_flag_result_variant_local_evaluation
  - test_get_feature_flag_result_boolean_decide
  - test_get_feature_flag_result_variant_decide

* Report combined errors when both errors_while_computing and flag_missing

When the server returns errorsWhileComputingFlags=true AND the requested
flag is not in the response, report both conditions as a comma-separated
string: "errors_while_computing_flags,flag_missing"

This provides better debugging context when both conditions occur.

* Add FeatureFlagError constants class for error type values

- Add FeatureFlagError class to types.py with constants:
  - ERRORS_WHILE_COMPUTING, FLAG_MISSING, QUOTA_LIMITED
  - TIMEOUT, CONNECTION_ERROR, UNKNOWN_ERROR
  - api_error(status) static method for dynamic error strings

- Update client.py to use FeatureFlagError constants instead of
  magic strings

- Update all tests to use constants for maintainability

This improves maintainability by:
- Single source of truth for error values
- IDE autocomplete and typo detection
- Documentation of analytics-stable values

* Remove print statements from test failure handlers

* Fix mypy type error in FeatureFlagError.api_error method

Accept Union[int, str] to match APIError.status type.
2025-12-15 17:06:24 -08:00
Dustin ByrneandGitHub af561acd9d feat: Add FlagDefinitionCacheProvider interface (#387)
* feat: Add FlagDefinitionCacheProvider interface

* feat: Add a Redis example for FlagDefinitionCacheProvider

* style: ruff format

* style: ruff format

* refactor: clean up mypy errors

* chore: mypy-baseline sync

* fix: Type Redis as Redis[str]

* fix: Adhere to strict typing

The defined types don't leave room for missing or optional keys. We'll
use the types as they're defined.
2025-12-12 11:31:55 -05:00
Tom PiccirelloandGitHub b8073984eb Publish to PyPI using Trusted Publisher (#388)
Twine [supports Trusted Publisher](https://github.com/pypa/twine/pull/1194/), but their documentation is a bit sparse.
2025-12-11 13:10:18 -08:00
Tom PiccirelloandGitHub 082a0e67cd Replace PAT with default GITHUB_TOKEN (#386)
* Replace PAT with default GITHUB_TOKEN

A PAT isn't needed for either of these Actions.

* Only expose env vars to step that needs them
2025-12-11 12:39:38 -08:00
Carlos MarchalandGitHub da9b730da1 feat(llma): multimodal-capture (#378) 2025-12-11 15:00:53 +00:00
Aleksander BłaszkiewiczandGitHub 381dc1c9ab fix: remove unused $exception_message and $exception_type (#383)
* fix: remove unused

* fix: remove exception type

* fix: wip
2025-12-10 11:33:14 +00:00
Dustin ByrneandGitHub e24b1038f3 feat: Allow customization of socket options (#385)
* feat: Allow customization of socket options

This allows clients to configure (e.g.) socket keep alive probes

* test: Prevent leaking modified _session

* chore: Add a type ignore

* feat: Add `disable_connection_reuse` config method

Disables connection pooling

* set_socket_options is idempotent
2025-12-09 01:52:24 -05:00
Aleksander BłaszkiewiczandGitHub f1b454b10d feat: mask values (#382)
* feat: mask values

* feat: wip

* fix: ruff

* feat: wip

* feat: wip

* fix: ruff

* fix: ruff

* feat: wip

* feat: wip

* feat: version bump
2025-12-05 19:35:56 +01:00
Dustin ByrneandGitHub 304218f78f fix: capture enriches with local eval when enabled (#380)
Captured events now use local evaluation results when
`send_feature_flags` is `True` and local evaluation is enabled.
2025-12-04 15:07:02 -05:00
Phil HaackandGitHub 671af32b02 feat(flags): Add ETag support for local evaluation polling (#381)
* Add ETag support for local evaluation polling

Add support for HTTP conditional requests using ETags to reduce bandwidth
when polling for feature flag definitions. When flag definitions haven't
changed, the server returns 304 Not Modified and the SDK skips processing.

- Add GetResponse dataclass to encapsulate response data, ETag, and status
- Update get() to send If-None-Match header and handle 304 responses
- Store ETag in client and pass it on subsequent polling requests
- Skip flag processing when 304 Not Modified is received

* Use _session rather than requests

Benefits:

1. Reuses TCP connections via keep-alive
2. 2 retries on connect/read errors
3. Faster handshakes

* Add unit tests for get() function

Test HTTP-level behavior including:
- ETag extraction from response headers
- If-None-Match header sent when etag provided
- 304 Not Modified response handling
- Fallback when 304 has no ETag header
- Error response handling (APIError)
- Authorization and User-Agent headers
- Timeout and URL construction

* Add defensive null check for response.data

Guard against unexpected None data in non-304 responses to prevent
TypeError when accessing dictionary keys.

* Clear stored ETag when server stops sending one

If the server stops including ETag headers in responses, clear the
stored ETag so we don't keep sending a stale If-None-Match header.

* Mask API tokens in log messages

Keep first 10 chars visible for identification while hiding the rest.
Addresses CodeQL security warning about logging sensitive data.

* Ran ruff format
2025-12-02 13:11:38 -08:00
Dylan MartinandGitHub 50a177e39a feat(flags): included evaluated_at properties in $feature_flag_called events (#374)
* format

* update tests

* bump version
2025-12-01 22:02:09 -05:00
Radu RaiceaandGitHub b7dfe68da7 feat(llma): add Gemini async (#375) 2025-11-27 11:35:54 -05:00
Carlos MarchalandGitHub b997548400 feat: add python 3.14 support (#373) 2025-11-25 16:43:34 +01:00
github-actions[bot] e54f8a3811 Update generated references 2025-11-15 12:44:50 +00:00
Aleksander BłaszkiewiczandGitHub e38bfb5e49 feat: use repr in code variables (#372) 2025-11-15 13:43:55 +01:00
Tue HaulundandGitHub b2c4043e09 fix: avoid overwriting consumer list when using more than one consumer (#370)
fix: avoid overwriting consumer list when using more than one consumer thread
2025-11-12 13:05:05 +01:00
github-actions[bot] f2f4fa6874 Update generated references 2025-11-11 18:13:36 +00:00
Carlos MarchalandGitHub 748553236d chore(llma): update SDKs (#367) 2025-11-11 18:12:33 +00:00
github-actions[bot] 15a9c1b861 Update generated references 2025-11-11 17:57:26 +00:00
Alessandro PogliaghiandGitHub 9972fa3f0b feat(ph-ai): PostHog properties dict in GenerationMetadata (#366) 2025-11-11 17:56:32 +00:00
github-actions[bot] 99422f02ab Update generated references 2025-11-11 08:57:51 +00:00
Carlos MarchalandGitHub 1dcfb99413 fix(llma): Langchain cache token double subtraction for non-Anthropic providers (#369) 2025-11-11 09:56:57 +01:00
Julian BezandGitHub 4303f78759 fix: add ruff check to CI and fix all lint errors (#360)
- Add 'Lint with ruff' step to CI workflow (was only running format check)
- Fix F401: Add explicit re-exports for public API types in __init__.py
- Fix F811: Rename duplicate test_openai_reasoning_tokens to test_openai_reasoning_tokens_o4_mini
- Fix E731: Convert lambda to def function in test_middleware.py
- Fix unused imports in client.py and test files (auto-fixed)

Without ruff check in CI, lint errors were accumulating on master undetected.
2025-11-10 12:45:52 +01:00
github-actions[bot] 1ec7afdc5f Update generated references 2025-11-07 15:57:55 +00:00
Aleksander BłaszkiewiczandGitHub 3a8876d1a9 fix: code variables without client (#368)
* fix: pass variables from init to client

* chore: version+changelog
2025-11-07 16:57:01 +01:00
Luke BeltonandGitHub a338772b84 fix docstring for set (#364) 2025-11-06 16:02:49 +00:00
github-actions[bot] 08bf769a6c Update generated references 2025-11-06 15:13:17 +00:00
Aleksander BłaszkiewiczandGitHub 8c64f43476 feat: add code variables capture (#365)
* feat: add code variables capture

* feat: bump version and add changelog
2025-11-06 16:12:17 +01:00
github-actions[bot] bb5f1b7130 Update generated references 2025-11-04 19:44:04 +00:00
Radu RaiceaandGitHub 34540a5d93 feat(llma): send number of web searches (#359)
* feat(llma): send number of web searches

* feat(llma): add more tests

* chore(llma): bump version

* fix(llma): feedback

* fix(llma): fix Gemini

* fix(llma): fix OpenAI's Chat Completions streaming
2025-11-04 19:43:11 +00:00
github-actions[bot] 0e502fa086 Update generated references 2025-11-03 12:29:46 +00:00
Julian BezandGitHub aa04cd6a7f fix(django): handle request.user in async middleware context (#358) 2025-11-03 12:28:57 +00:00
github-actions[bot] c141ab0c77 Update generated references 2025-11-02 18:58:55 +00:00
997f8228da fix(llma): cache cost calculation in the LangChain callback (#346)
* fix(llma): cache cost calculation in the LangChain callback

* fix: format

* Update posthog/ai/langchain/callbacks.py

Co-authored-by: Radu Raicea <radu@raicea.com>

* Bump version to 6.7.13

Master has already released 6.7.12 with other fixes, so this PR will be 6.7.13

---------

Co-authored-by: Radu Raicea <radu@raicea.com>
Co-authored-by: Andrew Maguire <andrewm4894@gmail.com>
2025-11-02 18:58:04 +00:00
github-actions[bot] 183ecc8edf Update generated references 2025-11-02 17:10:28 +00:00
Andrew MaguireandGitHub 2ff0810472 fix(llma): LangChain 1.0+ compatibility for CallbackHandler (#363)
* fix: Add LangChain 1.0+ compatibility for CallbackHandler imports

- Use try/except to import from langchain_core first (LangChain 1.0+)
- Fall back to legacy langchain imports for older versions
- Maintains backward compatibility with LangChain 0.x
- All existing tests pass (45 passed)

Fixes #362

* test: Add regression test for AgentAction/AgentFinish imports

- Tests that AgentAction and AgentFinish can be imported
- Tests on_agent_action and on_agent_finish callbacks with mock data
- Ensures compatibility with both LangChain 0.x and 1.0+
- Catches the import issue that was previously only tested with API keys

This addresses a test coverage gap identified during code review.

* chore: Add CHANGELOG entry for LangChain 1.0+ compatibility fix

* fix: Remove unused type: ignore comments for mypy

The type: ignore comments were only needed when the except block
executes, but CI runs with LangChain 1.0+ so the try block succeeds.
Mypy flags these as unused-ignore errors.

* chore: bump version to 6.7.12 for langchain 1.0 compatibility
2025-11-02 17:09:33 +00:00
Julian BezandGitHub 2cb5f02b1b fix(django): restore process_exception to capture view exceptions (#350)
Restores the process_exception method that was removed in v6.7.5 (PR #328),
which broke exception capture from Django views and downstream middleware.

Django converts view exceptions into responses before they propagate through
the middleware stack's __call__ method, so the context manager's exception
handler never sees them. Django provides these exceptions via the
process_exception hook instead.

Changes:
- Add process_exception method to capture exceptions from views and downstream
  middleware with proper request context and tags
- Add tests verifying process_exception behavior and settings (capture_exceptions,
  request_filter)
2025-10-29 10:40:08 +00:00
github-actions[bot] 450d96a5f1 Update generated references 2025-10-28 13:06:45 +00:00
Andrew MaguireandGitHub 5b4bab5fd7 chore: bump version to 6.7.11 for AI framework feature (#354)
Update version and changelog for PR #347
2025-10-28 13:05:50 +00:00
156170e34b feat(ai): Add $ai_framework property for framework integrations (#347)
* Add $ai_lib_metadata to AI integrations

Adds framework identification metadata to all AI events for easier filtering
and analytics. Each integration now includes a $ai_lib_metadata property with
schema version and framework name.

- LangChain: Hardcoded to "langchain"
- Native wrappers (Anthropic, OpenAI, Gemini): Uses provider name
- Ready for future frameworks (pydantic-ai, crewai, llamaindex)

This enables PostHog queries to easily distinguish between:
- Direct SDK wrapper usage
- Framework-mediated usage (LangChain, etc.)
- Different framework types

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

* Add \$ai_lib_metadata to sync/async paths and tests

- Added \$ai_lib_metadata to call_llm_and_track_usage (sync)
- Added \$ai_lib_metadata to call_llm_and_track_usage_async (async)
- Added test assertion in test_basic_completion
- Placed metadata at end of properties for consistency

All tests pass successfully.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

* Refactor: use unified utility function for $ai_lib_metadata

Creates a single `get_ai_lib_metadata(framework)` utility function to generate
the $ai_lib_metadata object, replacing inline implementations across the
codebase.

Changes:
- Add get_ai_lib_metadata() utility to utils.py
- Update LangChain callbacks to use utility function
- Update call_llm_and_track_usage() to use utility function
- Update call_llm_and_track_usage_async() to use utility function
- Update capture_streaming_event() to use utility function

Benefits:
- Consistency across all integrations
- Single source of truth for metadata structure
- Easier to extend with version detection later

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

* Add $ai_lib_metadata assertions to provider tests

Add missing $ai_lib_metadata assertions to Anthropic, Gemini, and LangChain tests to match the validation already present in OpenAI tests. Each test now verifies the metadata field contains the correct schema version and framework name.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

* Simplify to $ai_framework property, only for actual frameworks

Changes:
- Replace complex $ai_lib_metadata object with simple $ai_framework string
- Only include $ai_framework when using actual framework (LangChain)
- Remove $ai_framework from direct provider calls (OpenAI, Anthropic, Gemini)
- Update all tests to reflect new behavior

Before: {"schema": "v1", "frameworks": [{"name": "langchain"}]}
After: "langchain" (only when using LangChain framework)

This eliminates wasteful redundancy where framework=provider for direct calls.

---------

Co-authored-by: Claude <noreply@anthropic.com>
2025-10-28 09:59:20 +00:00
Julian BezandGitHub 141ced6402 chore: standardize workflow extensions to .yml (#349)
Rename workflow files from .yaml to .yml for consistency with existing
workflows (ci.yml, call-flags-project-board.yml).

This resolves naming confusion and standardizes all GitHub Actions
workflow files to use the .yml extension.
2025-10-24 14:17:27 +00:00
webjunkieandgithub-actions[bot] b1c193bcb1 Update generated references 2025-10-24 14:16:57 +00:00
Julian BezandGitHub 5e6f642650 fix(django): make middleware truly hybrid-compatible with sync and async Django stacks (#348)
Address code review feedback and critical issues from PR #328.

Changes:
- Keep __call__ as sync method that conditionally routes to __acall__ for async paths
- Use markcoroutinefunction() to properly mark instances when async is detected
- Detect async/sync at init time via iscoroutinefunction(get_response)
- Remove process_exception method - it was non-functional (Django doesn't call it on new-style middleware without MiddlewareMixin)
- Fix markcoroutinefunction fallback to be a simple no-op instead of accessing private API
- Exception capture works correctly via contexts.new_context() which has built-in exception handling
- Add comprehensive test coverage for sync, async, and hybrid middleware behavior
- Add async exception capture tests
- Refactor tests to use proper middleware initialization

This implementation follows Django's recommended hybrid middleware pattern where
both sync_capable and async_capable are True, allowing Django to pass requests
without conversion while the middleware adapts based on the detected mode.

The sync path behavior is identical to version 6.7.4 (pre-async), ensuring perfect
backward compatibility for WSGI deployments.

Addresses #329
Related to #328
2025-10-24 15:50:21 +02:00
Phil HaackandGitHub c892dc7a8a Bump version to 6.7.9 (#345) 2025-10-22 20:53:07 +00:00
Phil HaackandGitHub 52e023072e fix(flags): multi-condition flags with static cohorts returning wrong variants (#343)
* Fix multi-condition flags with static cohorts returning wrong variants

When a feature flag has multiple conditions and one contains a static
cohort, the SDK now correctly falls back to the API instead of
evaluating subsequent conditions locally and returning incorrect variants.

Introduce RequiresServerEvaluation exception to distinguish between:
- Missing server-side data (static cohorts) → immediate API fallback
- Evaluation errors (bad regex, missing properties) → try next condition

Changes:
- Add RequiresServerEvaluation exception class
- Update match_cohort() to throw RequiresServerEvaluation for static cohorts
- Update match_property_group() to propagate RequiresServerEvaluation
- Update match_feature_flag_properties() to handle both exception types
- Update client.py to catch both exceptions for API fallback
- Export RequiresServerEvaluation in __init__.py
- Add test for multi-condition static cohort scenario

All 84 feature flag tests pass.

* Add unit test for payloads

* ruff format
2025-10-21 13:40:36 -07:00
Radu RaiceaandGitHub 19d0562d41 fix(llma): missing await in OpenAI's streaming implementation (#342)
* fix(llma): missing async for OpenAI async

* chore(llma): bump version

* chore(llma): bump version
2025-10-16 14:46:26 +00:00
gewenyu99andgithub-actions[bot] ccbd9c23b7 Update generated references 2025-10-16 00:02:49 +00:00
Vincent (Wen Yu) GeandGitHub 71919f7125 fix: Check for references directory and generate 6.7.7 specs (#341)
* Fix check for directory and generate 6.7.7 specs

* Delete references
2025-10-15 20:02:16 -04:00
Vincent (Wen Yu) GeandGitHub 4d9a4b689e Run SDK generation after release job (#340)
* run after release job

* use bot pat

* Fix token placement

* Run this with a gh cli command
2025-10-15 19:45:33 -04:00
Manoel Aranda NetoandGitHub 890bb4b5f8 fix: remove deprecated attribute from exception events (#338) 2025-10-14 10:33:49 +00:00
Tom PiccirelloandGitHub 98a5bca68d chore: switch to fine-grained PAT (#337) 2025-10-13 10:19:12 -07:00
Vincent (Wen Yu) GeGitHubgreptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
898859b546 Add workflow to create and save versioned references (#332)
* Updates script to persist references

* Workflow to generate references to a folder

* Get rid of references, to be generated

* Update .github/workflows/generate-references.yaml

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

* Update .github/workflows/generate-references.yaml

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

* Review comments

* Update .github/workflows/generate-references.yaml

* Pin hashes and only run on releases

* Pin uv

---------

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
2025-09-30 17:02:57 -04:00
Carlos MarchalandGitHub c3ec12c4a6 feat: auto update llm sdks (#333)
* feat: auto update llm sdks

* fix: apply PR comments
2025-09-25 13:40:42 +02:00
Andy ZhaoandGitHub 65bf587fee fix: don't sort condition sets with variant overrides to the top (#330)
* fix: don't sort condition sets with variant overrides to the top

* fix test

* update test

* update version and change log
2025-09-22 14:10:43 -04:00
Dustin ByrneandGitHub 6b28b80b7a fix: Prevent core Client methods from raising exceptions (#327)
* fix: Prevent core Client methods from raising exceptions

The goal is to ensure that our client doesn't cause a panic in an
end-user application. This change updates
capture/set/set_once/group_identify/alias to swallow and log any
exceptions that occur. Note that this won't prevent errors from
propagating via the `on_error` callback if an error occurs while
processing the queue.

* test: Remove assertions that capture raises

These tests were broken anyways. Capture would only raise because it was
being called with no arguments, not because api_key or host are None.
2025-09-17 15:47:07 -04:00
Oliver BrowneandGitHub 466bc358b2 feat(err): add __acall__ to django middleware (#328)
* add __acall__

* fix types
2025-09-16 15:40:11 +03:00
Carlos MarchalandGitHub 2508f0193d fix/system prompt sometimes missing (#326)
* fix: always capture system prompt

* chore: bump version

* fix: gemini system prompt capture

* chore: imports at top

* fix: test

The mock we were passing from this test
reporetd that it had a `system instruction` field,
breaking assumptions

* chore: lint

* fix: better code organization

* chore: lint
2025-09-05 17:28:55 +02:00
Radu RaiceaandGitHub 7e6009bc3e fix(llma): gemini missing cached and reasoning tokens (#323)
* fix(llma): Gemini missing cached and reasoning tokens

* chore(llma): bump version

* chore(llma): run ruff
2025-09-04 14:21:32 -04:00
Radu RaiceaandGitHub 63d8c8dc0f fix(llma): streaming providers with tool calls (#319)
* fix(llma): tool calls in streaming Anthropic

* fix(llma): Gemini content

* fix(llma): extract converters for providers

* fix(llma): continuation of DRY refactoring

* fix(llma): add $ai_tools to streaming Gemini

* fix(llma): tool calls in streaming Gemini

* fix(llma): tool calls in streaming OpenAI Chat Completions

* fix(llma): fix test

* fix(llma): run ruff

* fix(llma): fix types

* fix(llma): run ruff

* chore(llma): run mypy baseline sync

* chore(llma): bump version

* fix(llma): fix test

* chore(llma): update CHANGELOG

* fix(llma): Responses API streaming tokens

* fix(llma): run ruff

* fix(llma): run ruff
2025-09-03 20:02:38 +00:00
Dylan MartinandGitHub 2d418d1224 technically incorrect (#321) 2025-09-02 17:02:41 -07:00
Carlos MarchalandGitHub a70ed0506e feat(llma): redact base64 images (#318) 2025-09-01 09:13:07 +02:00
Phil HaackandGitHub b7817791aa Bump version to 6.7.0 (#317) 2025-08-26 22:32:04 +00:00
Phil HaackandGitHub b7c71b09aa fix(flags): flag dependency evaluation for multivariate flags (#316) 2025-08-25 14:20:45 -07:00
Vincent (Wen Yu) GeGitHubgreptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
4dacaed27b Add categories to doc specs (#313)
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
2025-08-21 11:48:54 -04:00
Juraj MajerikandGitHub 6947dc6635 chore: bump version to 6.6.1 (#314) 2025-08-21 16:14:27 +02:00
Juraj MajerikandGitHub 44dddb1a4a Fix NoneType error when group_properties is None (#312) 2025-08-19 12:01:12 -07:00
Phil HaackandGitHub 413fb5af77 feat(flags): implement local evaluation of flag dependency filters (#311) 2025-08-19 09:43:46 -07:00
Phil HaackGitHubCopilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>
7f087e25b3 chore: Improvements to example.py (#310)
Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>
2025-08-18 21:46:58 -07:00
Dylan MartinandGitHub ae38bbabc1 feat(flags): support passing in lists of flag keys to the /flags endpoint instead of evaluating every flag every time we fall back (#307) 2025-08-18 16:10:34 -07:00
David NewellandGitHub cb75671cd9 feat: we should capture which properties were added as tags (#304) 2025-08-08 11:44:34 +01:00
Vincent (Wen Yu) GeandGitHub edc0d1a1ed Remove placeholder for params, waste space (#298) 2025-08-07 15:20:27 -04:00
Phil HaackandGitHub 82a7f7c171 fix(flags): Pass project API key in remote_config requests (#303) 2025-08-06 21:17:30 +00:00
Radu RaiceaandGitHub 698f79b257 feat(llmo): support Vertex AI (#302)
* feat(llmo): support Vertex AI

* chore(llmo): run formatter

* fix(llmo): fix types error

* chore(llmo): run formatter

* chore(llmo): bump version
2025-08-05 15:33:10 -04:00
Radu RaiceaandGitHub 947ee46661 fix(llmo): tool calls are broken for most providers (#299)
* fix(llmo): set the $ai_tools properly for all providers

* fix(llmo): remove privacy mode from $ai_tools

* chore(llmo): bump version

* chore(llmo): run formatter

* fix(llmo): properly set tool calls in $ai_output_choices

* chore(llmo): bump version

* chore(llmo): run formatter

* fix(llmo): fix types error

* feat(llmo): change $ai_output_choices to have an array of content

* chore(llmo): run formatter

* feat(llmo): create text type object

* chore(llmo): update CHANGELOG.md
2025-08-05 14:02:37 -04:00
Phil HaackGitHubgreptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
d5bac13022 feat: Expose get_feature_flag_result method in public API (#284)
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
2025-08-05 10:14:29 -07:00
Dylan MartinandGitHub f73774f616 fix (#300) 2025-08-01 17:32:58 -07:00
Radu RaiceaandGitHub 6ecd4f7efb bug(llmo): fix anthropic tool call response (#297)
* bug(llmo): fix anthropic's tool call response

* bug(llmo): fix tool calls response handling for anthropic

* bug(llmo): run formatter

* bug(llmo): bump version

* bug(llmo): add date to changelog
2025-07-31 17:13:44 -04:00
Georgiy TarasovandGitHub 864d718412 fix(ai): capture tool calls in reasoning models (#292)
* fix: capture tool calls in reasoning models

* fix: check for empty tool calls
2025-07-23 11:51:58 +02:00
Dylan MartinandGitHub 5a5326b66a feat(flags): make the sendFeatureFlags parameter more declarative and ergonomic (#283) 2025-07-22 15:20:54 -07:00
Radu RaiceaandGitHub 3289cfe3af feat(llmo): Use default PH client for LangChain (#293)
* feat(llmo): Use default PH client for langchain

* chore: Run formatter

* feat: Test the CallbackHandler without any PH client

* chore: Run formatter
2025-07-22 14:09:47 -07:00
2fb3c980a5 feat(llmo): Make it optional to pass posthog client (#291)
Co-authored-by: Peter Kirkham <peter@posthog.com>
2025-07-22 06:53:47 +00:00
David NewellandGitHub f8efe050f2 fix: capture django processed exceptions (#287) 2025-07-16 21:54:01 +02:00
Vincent (Wen Yu) GeandGitHub f600b0e7de Init reference doc generation (#280) 2025-07-15 13:53:13 -04:00
Phil HaackandGitHub 4ef129d646 fix: Ignore new flag filter type in local evaluation (#285) 2025-07-11 16:33:23 +00:00
Dylan MartinandGitHub 633504ea56 feat(flags): decouple local evaluation from personal API keys; support decrypting remote config payloads without relying on the feature flags poller (#282) 2025-07-10 08:03:37 -07:00
Oliver BrowneandGitHub 1830d7917d fix: add middleware setting for custom client (#281)
* Add middleware setting for custom client

* mypy

* comment
2025-07-09 17:06:20 +03:00
Dylan MartinandGitHub 70da872508 feat(flags): add a flag_fallback_cache that tracks feature flag evaluation results and uses them as fallback values whenever the /flags API isn't available (#275) 2025-07-07 07:13:35 +00:00
Oliver BrowneandGitHub ae0a4e7b57 fix: default send_feature_flags false for capture_exception (#278)
* default send_feature_flags false

* bump version
2025-07-02 22:21:25 +03:00
Oliver BrowneandGitHub 66bf50c6ac fix(err): permit disabling person processing (#277)
* whoops

* bump version
2025-07-01 15:26:57 +03:00
Oliver BrowneGitHubgreptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>David Newell
5f5b68cbbb feat: prep for 6.0.0, bunch of breaking changes (#273)
* alright

* fix exports

* Update posthog/test/test_before_send.py

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

* Update posthog/__init__.py

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

* Update posthog/__init__.py

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

* fix tests after rebase

* update examples

* update mypy baseline

* Revert "update mypy baseline"

This reverts commit da395dd7cc075d1f5b1c03d748544e4abcd75bb9.

* try again

* whatever

* request user

* fix middleware

* getattr is_authenticated

* allow using custom client for exception capture

* update comments

* fix circular import

* type arguments, use TypedDict

* fix setup

* mypy

* whoops

* ok

* further mypy

* mypy sync

* docs and fixes

---------

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Co-authored-by: David Newell <david@posthog.com>
2025-06-27 14:57:02 +03:00
Lucas RicoyandGitHub aa494972b0 docs: update release details on readme (#272) 2025-06-21 00:38:18 +01:00
Lucas RicoyGitHubgreptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
59174b81b5 feat: bump version to 5.4.0 with session_id on page method (#271)
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
2025-06-21 00:18:57 +01:00
Lucas RicoyandGitHub 25b221cc2e chore: ensure session_id context works with page method (#269) 2025-06-20 23:49:53 +01:00
Oliver BrowneandGitHub 18f3f89759 fix(err): always safe_str exception values (#267)
* always safe_str

* bump version
2025-06-19 14:38:17 +00:00
Oliver BrowneandGitHub a13acafb81 feat(err): construct full trace if no traceback available (#266)
* construct full trace if no traceback available

* delete asserts

* bump version
2025-06-19 16:47:42 +03:00
Oliver BrowneandGitHub eb3b524d72 document prep_local (#265) 2025-06-18 19:50:32 +01:00
Oliver BrowneandGitHub d5da699816 feat: session and identity integrate with context now (#264)
* session and identity in context

* bump version

* make django integration use context distinct id and session functions

* don't use self

* fix comments

* fix middleware tests

* clarify fresh and distinct id's

* Fix exported modules, add makefile command to test changes locally

* tiny fix
2025-06-18 19:27:27 +03:00
b521c2acbd feat(err): add django middleware (#263)
* fix exactly-once capture

* add middleware

* fix typing

* ignore unreacable

* Revert "ignore unreacable"

This reverts commit 0458f0efa6c8e52ecfb1eeb41c57a76d84578164.

* add unreachable ignore

* move unreachable ignore

* switch to use request.headers

* clarify comment

* Update posthog/integrations/django.py

Co-authored-by: David Newell <d.newell1@outlook.com>

* explain typle

* fix comment

* explain that tags become properties

* fix tests

---------

Co-authored-by: David Newell <d.newell1@outlook.com>
2025-06-17 17:54:32 +03:00
Oliver BrowneandGitHub fe0a680e3f fix: delete sentry integration (#262)
* delete relevant files

* bump uv lock, bump major version as deprecation

* README.md

* try mypy sync

* Revert "try mypy sync"

This reverts commit e1b98b26e59132e52eff6389afd42cb1b07a6a0b.

* try looking at the github action
2025-06-16 18:38:53 +03:00
Oliver BrowneandGitHub eba2b112dc fix(err): just check if the passed exceptions is a BaseException (#261)
* just check if it's an exception first

* version bump
2025-06-13 19:36:00 +00:00
Georgiy TarasovandGitHub 5c2897a767 feat(ai): langchain cached and reasoning tokens (#258)
* fix: reasoning and cached tokens

* test: new flows

* fix: missing field

* chore: bump

* fix: make sure we send write/read/reasoning tokens
2025-06-13 15:02:06 +02:00
Phil HaackandGitHub c85e361642 Add flags project board workflow (#259) 2025-06-12 16:33:42 +00:00
Oliver BrowneandGitHub c2afc4d27a fix: change scoped export, add capturing param (#257)
* change export, add capturing param

* capturing -> capture_exceptions
2025-06-12 10:29:32 +01:00
Peter KirkhamandGitHub 4625267f16 feat: add support for parse via responses (#256) 2025-06-11 05:39:24 +01:00
David NewellandGitHub 2e4ec31e8e remove 'import posthog' (#255) 2025-06-10 11:05:22 +01:00
David NewellandGitHub 5b853fe9a1 fix: makefile import (#254) 2025-06-09 19:20:31 +01:00
David NewellandGitHub 4521db34a1 feat: more django context (#252) 2025-06-09 14:32:42 +01:00
Paul D'AmbraandGitHub 5c84dcfccf feat: add before_send callback (#249) 2025-06-09 13:56:46 +01:00
Paul D'AmbraandGitHub c4ca591227 fix changelog to match pypi (#253) 2025-06-09 12:09:15 +01:00
Paul D'AmbraandGitHub 00c476328c chore: more fiddling to get release working (#251)
* chore: more fiddling to get release working

* fix

* fix
2025-06-09 11:50:30 +01:00
Paul D'AmbraandGitHub 7e018382ab chore: use uv run when releasing (#250) 2025-06-09 10:11:12 +00:00
Paul D'AmbraandGitHub 312140ff25 fix: makefile for posthog_analytics release (#248) 2025-06-09 10:51:04 +01:00
Paul D'AmbraandGitHub 7db912f06f chore: pyproject and CI update (#247) 2025-06-07 14:08:47 +03:00
Dylan MartinandGitHub 3b07b6ed47 chore(flags): roll everyone onto /flags (#246) 2025-06-06 16:36:10 -07:00
Oliver BrowneandGitHub 8b136d433b fix: lets try again (#244)
* maybe

* bump version
2025-06-06 16:25:47 +02:00
ce2a9bebc6 fix: python release attempt 3 (#242)
* fix: maybe the classifier is deprecated

* fix

* fix

* fix

* fix

* fix

* fix

* fix

* fix

* fix

* fix

* ruff

* bump version for release

---------

Co-authored-by: Oliver Browne <oliver@posthog.com>
2025-06-06 15:16:31 +03:00
Oliver BrowneandGitHub 6fb9e6d4c0 fix: bump for release (#243)
* bump for release

* changelog

* changelog
2025-06-06 11:59:58 +00:00
Oliver BrowneGitHubgreptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>David Newell
c62e60b9d3 feat(err): add context manager and tag functions (#239)
* add context maanager and tag functions

* Update posthog/scopes.py

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

* ran black

* black locally disagrees with ci. python is awful

* bleh

* isort

* fix mypyp thing

* Revert "fix mypyp thing"

This reverts commit 21ad8733610967cad0bbf8451508ca1189be543c.

* update baseline

* lets try again

* alright lets try again

* revert to baseline

* try ignoring it i guess

* black

* try supporting async too

* black

* mypy

* formatting

* fix changelog

* fix comment

* we only support python 3.9+

* change decorator name

* fix example

* isort

* auto-capture in with blocks

* fix tests

* black

* inherit tags by default

* assert swap

* add tags to all events

* rm comment

* Update example.py

Co-authored-by: David Newell <d.newell1@outlook.com>

---------

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Co-authored-by: David Newell <d.newell1@outlook.com>
2025-06-06 14:49:15 +03:00
Paul D'AmbraandGitHub 93d10d4d10 fix: release action failed (#241) 2025-06-05 15:50:27 +01:00
Paul D'AmbraandGitHub a03e2a066f fix: opinionated setup and clean fn fix (#240) 2025-06-05 14:25:56 +01:00
fc8812806c pin actions versions (#210)
* pin actions versions
---------

Co-authored-by: Paweł Szczur <orian@users.noreply.github.com>
2025-05-27 08:25:50 +00:00
Paweł SzczurandGitHub 9b43c0544d fix: feature flag request use geoip_disable (#235)
* make feature flag request use geoip_disable
2025-05-27 10:18:55 +02:00
Peter KirkhamandGitHub ed12f1b208 feat: support gemini (#237) 2025-05-24 00:23:34 +01:00
Peter KirkhamandGitHub af6882ea00 feat: composition over inheritance (#236) 2025-05-23 01:34:04 +01:00
Rafael AudibertandGitHub 8253c66412 feat: Remove deprecated monotonic lib (#231) 2025-04-29 11:14:59 -03:00
Rafael AudibertandGitHub 11461049fd chore: Release automatically when changed version.py (#232) 2025-04-29 11:14:50 -03:00
Phil HaackGitHubgreptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
74ea024265 feat: Add new FeatureFlagResult class and tests (#227)
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
2025-04-24 10:09:29 -07:00
Phil HaackandGitHub 265de15b6b chore: Add parameters to bin/test (#228) 2025-04-23 13:53:28 -07:00
Phil HaackGitHubgreptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
62618b7547 chore: Make condition_index optional. Also added some scripts for local dev. (#223)
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
2025-04-18 11:59:49 -07:00
Dylan MartinandGitHub 849a81decf chore(flags): latest version of posthog-python now uses /flags by default, except for a few exceptions (#222)
* wahahaha

* fix tests

* whoops don't forget to roll it out
2025-04-15 17:13:26 -04:00
Dylan MartinandGitHub c08d05d894 fix(flags): pass in the correct hashes. (#221)
* shoot

* cut customer token

* bump version

* dump posthog api from excluded
2025-04-15 16:09:47 -04:00
Dylan MartinandGitHub c07e7c87ee ubunut 20-04 is EOL (#220) 2025-04-15 13:44:19 -04:00
Dylan MartinandGitHub 0a3c06086a chore(flags): roll 10% of posthog-python /decide traffic (and all of PostHog's personal SDK traffic) to /flags (#218)
* init

* moved the constants

* formatting

* mypy

* fr do some damn formatting

* don't exclude posthog

* make it a set

* differentiate

* fix AI test

* use the same type everywhere

* ready to release
2025-04-15 12:57:49 -04:00
8ce657630f test(llm-observability): Account for LangGraph 0.3.29 changes (#219)
* test(llm-observability): Account for LangGraph 0.3.29 changes

* formatting

---------

Co-authored-by: dylan <dylan@posthog.com>
2025-04-15 12:36:15 -04:00
David NewellandGitHub 28ef3c7e8e fix: add field to proxy client setup (#217) 2025-04-11 13:57:26 +01:00
David NewellandGitHub 00d2f9066f chore: update license (#213) 2025-04-10 15:30:54 +01:00
David NewellandGitHub 3a07d907f7 feat: log captured exceptions (#215) 2025-04-10 12:14:29 +01:00
David NewellandGitHub bd08b51ac8 fix failing ai test (#216) 2025-04-10 12:01:25 +01:00
Frank HamandandGitHub a5c738f2cb bump version (#209) 2025-03-26 16:10:26 +00:00
Frank HamandandGitHub e337ff51a9 update automatic retries to include read errors (#208)
in e.g. lambda environments the connection can time out between invocations,
this comes through to the client as a "RemoteDisconnected" error, which it
turns out urllib classifies as a "read" error not a connection error (as
it's possible to get this error after data has been sent)
2025-03-26 16:02:08 +00:00
Phil HaackandGitHub 4ab73eeac2 feat(flags): Add more details such as version, id, and reason to $feature_flag_called events (#207)
* Flesh out Decide response types

* Ensure we normalize get_decide

In a back compat manner.

* Populate feature_flags_by_key when setting feature_flags

Since `self.feature_flags_by_key` is derived from `self.feature_flags`, and we often set the latter in unit tests, but forget to set the former, our tests can be wonky.

This ensures that when we set `self.feature_flags`, we always set `self. feature_flags_by_key`

* Annotate types

* Lookup local flag by key

Fixes #121

* Refactor local flag evaluation into its own method

* Include extra details in `$feature_flag_called` events

* Fix up type annotations, tests, and formatting

* Update lib to decide v4

* Bump version and add changelog
2025-03-25 16:27:24 -07:00
RossandGitHub 7635457f90 feat: Support serializing dataclasses (#206)
* Support serializing dataclasses

* Update version

* Run black

* Fix for Python 3.9
2025-03-17 14:28:33 +00:00
Peter KirkhamandGitHub 041adad7fb feat: add support for responses api (#205)
* feat: add suppoort for responses api

* fix: test

* fix: black

* fix: test - hopefully

* fix: test - hopefully #2

* fix: test - hopefully #3

* fix: test - hopefully #4

* fix: greptaile catch

* fix: mypy is not my friend

* fix: isort usort weallsort

* fix: noredef

* fix: mypy baseline

* fix: mypy

* fix: mypy
2025-03-14 05:16:52 +00:00
David NewellandGitHub 2c972ab7bc fix: distro requirement for analytics package (#204) 2025-03-12 14:12:15 +00:00
Peter KirkhamandGitHub 7b16bcb62e fix: azure open ai delta check (#203) 2025-03-10 21:35:36 +00:00
Peter KirkhamandGitHub dc765e6ee1 feat: cached tokens (#202)
* feat: cached tokens

* feat: add tool support

* chore: local test

* chore: isort black

* chore: bump v

* chore: remove import

* fix: types

* fix: black

* fix: mypy unpacking of None

* chore: mypy baseline

* feat: mypy fix

* fix: did things and stuff

* fix: mypy yourpy whos py?

* fix: things can be None

* fix: move test

* fix remove exampels from package

* fix: losing my py
2025-03-06 22:37:21 +00:00
Dylan MartinandGitHub f36daa199f chore(flags): improved some logs for quota limiting (#197)
* haha okay

* tests workin

* format

* use case-sensitive comparisons

* omg LOL

* fix tests

* jeez

* this will probably work

* now do local eval

* okay

* yo

* formatting

* fix import order

* type check

* ai yi yi

* code review

* format

* do it

* merge conflict UGH

* black formatting

* bump version

* correct changelog
2025-03-03 14:00:52 -05:00
Peter KirkhamGitHubgreptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
44a8637e5a feat: azure export w/ async (#200)
* feat: azure export w/ async

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
2025-02-28 20:40:09 +00:00
Michael MatlokaandGitHub f874fe53cb feat(llm-observability): $ai_tools capture in Langchain (#199) 2025-02-27 17:50:20 +00:00
Paul D'AmbraandGitHub f84ad25963 feat: add some platform info to events (#198) 2025-02-26 12:26:17 +00:00
Peter KirkhamandGitHub f05e716224 fix: async client (#196) 2025-02-23 13:10:43 +00:00
Dylan MartinandGitHub a74b8a7fcc feat(feature-flags): support quota limiting for feature flags (#195)
* haha okay

* tests workin

* format

* use case-sensitive comparisons

* omg LOL

* fix tests

* jeez

* this will probably work

* now do local eval

* okay

* yo

* formatting

* fix import order

* type check

* ai yi yi

* code review

* format
2025-02-21 15:45:51 -05:00
Dylan MartinandGitHub 38eba50f21 fix(flags): remove lower() when evaluating feature flag payloads – these payloads are case-sensitive! (#191)
* haha okay

* tests workin

* format

* use case-sensitive comparisons

* omg LOL

* fix tests

* jeez
2025-02-19 19:51:40 -05:00
Peter KirkhamandGitHub beac1539dc fix: support usage as part of generation (#192) 2025-02-18 00:17:52 +00:00
HavenGitHubgreptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>Manoel Aranda Neto
0375c7eb7f feat(flags): Add method for fetching decrypted remote config flag payload (#180)
* feat(flags): Add method for fetching decrypted remote config flag payload

* tweak

* tweak

* tweak

* Update posthog/__init__.py

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

* tweak

* get example script working

* format

* sort import

* tweak

* bump minor version

* Update posthog/version.py

Co-authored-by: Manoel Aranda Neto <5731772+marandaneto@users.noreply.github.com>

* Use flag key instead of id

* tweak

* tweak

---------

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Co-authored-by: Manoel Aranda Neto <5731772+marandaneto@users.noreply.github.com>
2025-02-13 14:21:05 -08:00
Frank HamandandGitHub 9705b53da6 automatically retry connection errors (#190)
* automatically retry connection errors

from the docs for max_retries: this applies only to failed DNS lookups,
 socket connections and connection timeouts

* run tests on multiple python versions

* update freezegun
2025-02-12 12:29:01 +00:00
Phil HaackandGitHub b91c5acc95 Add mypy to CI (#189) 2025-02-11 09:50:45 -08:00
Phil HaackandGitHub 50e1df93a0 Use casefold to compare strings case insensitively (#184) 2025-02-11 08:06:20 -08:00
Rafael AudibertandGitHub 0634ee8cbd fix: Move code under mypy type (#188)
* fix: Move code under mypy type

This is incorrect, we should've added these slightly lower in the method definition to avoid mypy from breaking

* feat: Bump to 3.12.1
2025-02-11 12:01:08 -03:00
Phil HaackandGitHub 99c72f6ef0 Remove the usage of is_simple_flag (#186) 2025-02-10 18:52:54 -08:00
Phil HaackandGitHub 30b8ec24e7 Stop capturing all feature flags on $feature_flag_called event. (#181) 2025-02-10 17:39:03 -08:00
Peter KirkhamandGitHub 0692b846af feat: add beta parse method support (#185) 2025-02-11 00:43:36 +00:00
Phil HaackandGitHub 49de45888e Deprecate the context argument (#182) 2025-02-10 15:26:29 -08:00
Phil HaackandGitHub 53de042ffb Allow specifying the flag in the example script (#157)
* Allow specifying the flag in the example script

* Reformat

* Run isort
2025-02-07 09:28:00 +09:00
Phil HaackandGitHub 7c7287f303 Move accessing variants outside of loop (#175)
* Move accessing variants outside of loop

`flag_variants` doesn't depend on condition so it doesn't make sense to declare it in the loop.

* Fix assertion

* Remove incorrect comment

Comment seems superfluous anyways.

* Break out of the loop when the key is found

The purpose of the loop is to loop through the flag keys and evaluate the one where `flag["key"] == key`. Once that key is found, there's no need to continue the loop.

* Complete the test

Looks like the test was missing an assert.

* Precompute valid variant keys outside loop
2025-02-07 09:24:17 +09:00
Peter KirkhamandGitHub 198aecf9c7 fix: langchain tool parent add (#179) 2025-02-05 19:03:14 +00:00
277868e12a feat(llm-observability): LangChain spans (#176)
* feat: refactor to dataclasses

* feat: spans

* test: fix part 1

* test: fix part 2

* test: fix part 3

* test: fix part n

* test: add langgraph agent test

* chore: bump and linters

* chore: bump

* fix: correctly capture a parent id when a custom trace_id is set

* test: multiple spans parent_ids

* fix: exception serialization

* refactor: ai_trace_name -> ai_span_name and ai_generation_id -> ai_span_id

* fix: logs typos

* Add minor breaking change note to changelog

* fix: naming

---------

Co-authored-by: Michael Matloka <michael@matloka.com>
2025-01-28 14:23:22 +01:00
Georgiy TarasovandGitHub 833833bfbc fix: serialize pydantic models (#177) 2025-01-27 17:38:49 +00:00
Peter KirkhamandGitHub 999aead498 feat(ai): add error handling to python ai sdk (#174) 2025-01-24 21:09:49 +00:00
Georgiy TarasovandGitHub 07dbf46fbb fix(llm-observability): parallel traces (#172)
* fix: parallel traces

* fix: linters

* chore: bump

* fix: better naming for clarity
2025-01-23 17:27:46 +01:00
Michael Matloka 4aa32ee6fa fix: Actually fix LangChain callback in posthoganalytics 2025-01-22 16:27:36 +01:00
Michael MatlokaandGitHub 869f87af26 fix: Account for import posthog in posthoganalytics release (#171) 2025-01-22 13:50:39 +00:00
Peter KirkhamandGitHub 9a12b71dcc chore: bump (#170) 2025-01-21 23:33:47 +00:00
Michael MatlokaandGitHub 8e05c1951c feat(llm-observability): LangChain tracing, with LangGraph tests (#169) 2025-01-21 23:18:55 +00:00
91ac7bc82f fix(llm-observability): capture system prompt for anthropic (#167)
Co-authored-by: Peter Kirkham <peter@posthog.com>
2025-01-17 21:04:37 +00:00
ef0b92d9ff Feat: Add Anthropic to Python SDK (#165)
Co-authored-by: Georgiy Tarasov <gtarasov.work@gmail.com>
2025-01-17 20:33:48 +00:00
Georgiy TarasovandGitHub 267ab0440b fix(llm-observability): flatten langchain's additional_kwargs (#166)
* fix: flatten additional_kwargs

* fix: remove print
2025-01-17 17:59:38 +01:00
Peter KirkhamandGitHub cd510cc386 feat: privacy_mode (#164) 2025-01-15 01:28:52 +00:00
Georgiy TarasovandGitHub c377be7140 feat(llm-observability): add new packages for posthoganalytics (#163) 2025-01-14 10:50:46 +01:00
Georgiy TarasovandGitHub 113de04672 fix(llm-observability): include the ai packages (#162)
* fix: setuptools

* fix: include packages
2025-01-14 10:27:05 +01:00
Peter KirkhamandGitHub 2adef77a18 fix: update to export module (#161) 2025-01-14 01:25:00 +00:00
Peter KirkhamandGitHub 211075dabe feat: Embeddings + Personless events + Destructure property JSON (#160) 2025-01-14 00:35:09 +00:00
Georgiy TarasovandGitHub a38360d087 feat(llm-observability): add langchain integration (#159)
* feat(ai): LangChain integration v0.1

* test: langchain integration tests

* test: langchain-openai for v2 and v3

* chore: reorganize imports

* fix: ci

* fix: set python on ci to 3.9

* fix: upgrade ci for python 3.9

* fix: fallback for distinct_id

* fix: personless events for omitted distinct_ids

* fix: review comments

* feat: base url retrieval
2025-01-13 18:40:02 +01:00
e6e4368852 Feat: Add llm observability to python sdk (#158)
Co-authored-by: Michael Matloka <michael@matloka.com>
2025-01-11 01:34:27 +00:00
Sibin M SandGitHub 93623e8a5a [FEATURE]Add distinct_id to group_identify (#155)
* [FEATURE]Add distinct_id to group_identify

* [TESTS]Updated test cases for adding distinct_id to group_identify

* [LINT-FIX]client.py and test_client.py

* [CHORE]Verion bump and changelog update
2025-01-03 16:00:35 -05:00
Dylan MartinandGitHub bead19a9c2 fix: CI (#156)
* test CI

* heck it, upgrade python

* okay don't do anything silly with the cache hits i guess

* more CI upgrades :crossedfingers

* upgrade all CI to latest versions, then

* jk this is how python works

* whackamole

* what even

* yeesh

* this can't be it

* if this breaks ill kms

* dark magic dark MAGIC

* im giving up on my dreams
2025-01-03 15:50:21 -05:00
Dylan MartinandGitHub cc0568b0a6 Update CODEOWNERS (#154) 2025-01-02 12:52:51 -05:00
Dylan MartinandGitHub ddd744b51a fix(flags): bump the version (#148) 2024-11-27 17:15:38 -05:00
5c4cc35bfe fix(flags): correctly emit feature flag events with the FF response on get_feature_flag_payload calls (#143)
* this is the fix, needs tests

* fix test

* tests

* yeah

* please work

* ran the formatter

* code review feedback

* how'd this get here

* bump version add changelog

* Update CHANGELOG.md

Co-authored-by: David Newell <d.newell1@outlook.com>

---------

Co-authored-by: David Newell <d.newell1@outlook.com>
2024-11-25 14:51:06 -05:00
cde24f1525 fix(err): fix distinct_id, set personless and use a uuid (#144)
Co-authored-by: David Newell <david@posthog.com>
2024-11-25 12:09:57 +00:00
David NewellandGitHub fdc154deb9 chore: add type to stack (#142) 2024-11-19 12:46:52 +00:00
David NewellandGitHub 65abf5f34a chore: add lang to frames (#139) 2024-10-24 16:18:02 +01:00
David NewellandGitHub e4358be3d2 feat: add super properties (#138) 2024-10-03 17:07:47 +01:00
5e73f7b018 Remove deprecated datetime.utcnow() in favour of datetime.now(tz=tzutc()) (#124)
Co-authored-by: Neil Kakkar <neilkakkar@gmail.com>
2024-09-24 11:09:15 +01:00
Neil KakkarandGitHub 2898d08d8f fix(errors): Make sure project root exists to judge in app frames (#136)
* fix(errors): Make sure project root exists to judge in app frames

* prep release
2024-09-16 11:11:28 +01:00
Neil KakkarandGitHub 6afa42bb2e fix: Update Django integration for manual capture (#135) 2024-09-10 09:27:47 +01:00
Neil KakkarandGitHub 9886e6fbaa feat(errors): Add manual exception capture (#134)
* feat(errors): Add manual exception capture

* prep release

* use backwards compatible helper

* add tests
2024-09-09 11:40:51 +01:00
Neil KakkarandGitHub b468a9064a fix(setup): Make sure all packages are bundled (#133)
* fix(setup): Make sure all packages are bundled

* prep release

* black
2024-09-03 08:13:10 +01:00
Neil KakkarandGitHub 7576234533 fix(setup): Make sure all packages are bundled (#132)
* fix(setup): Make sure all packages are bundled

* prep release
2024-09-03 07:23:54 +01:00
Neil KakkarandGitHub 7cd12b1421 feat(errors): Add django integration and in app frames (#131) 2024-09-03 06:10:36 +01:00
Neil KakkarandGitHub 749f7b5940 feat(error-capture): Add basic exception autocapture (#128) 2024-08-28 06:59:41 +01:00
Phani RajandGitHub 98e82f9ae6 bump version to 3.5.2 (#130) 2024-08-21 11:58:02 -05:00
Phani RajandGitHub 32ad34f8b1 fix(feature flags): Guard for None values when comparing person Properties (#129)
* Guard for None values when comparing person Properties
2024-08-21 11:18:20 -05:00
Frank HamandandGitHub e0ae599d97 Switch us-api.i hosts to just us.i (#119)
We dropped the -api as it's a bit confusing (is capture really api?)
2024-08-14 08:41:10 +01:00
Neil KakkarandGitHub 07396e64a2 Create CODEOWNERS (#122) 2024-05-02 18:02:13 +01:00
Neil KakkarandGitHub 0e264a8cb2 fix(flags): Add a shorter configurable timeout for flag requests (#120) 2024-03-04 14:25:11 +00:00
Brett HoernerandGitHub 6785720db1 Bump to 3.4.2 (#118) 2024-02-20 08:43:10 -07:00
Brett HoernerandGitHub 97955fe0c9 Add historical_migration option to toplevel Client (#117) 2024-02-15 06:42:43 -07:00
Neil KakkarandGitHub 1b21afe570 chore: Use ingestion hosts for event capture (#116) 2024-02-13 11:19:21 +00:00
Neil KakkarandGitHub af424c833d fix(routing): Update hosts to point to right ingestion host (#115) 2024-02-05 12:14:14 +00:00
Neil KakkarandGitHub 40b5bea6c6 chore(flags): Update type hints for newer mypy versions (#114)
* chore(flags): Update type hints for newer mypy versions

* black
2024-01-30 16:06:10 +00:00
Neil KakkarandGitHub 230d9c2dbb fix(flags): Update relative date op names (#113) 2024-01-26 15:52:46 +00:00
635db706e8 Module functions to also return the same as its Client equivalent (#111)
Co-authored-by: Neil Kakkar <neilkakkar@gmail.com>
2024-01-19 11:13:52 +00:00
Neil KakkarandGitHub 4b9667ccb3 fix(flags): Don't override existing props when adding flags (#110) 2024-01-10 14:34:23 +00:00
Neil KakkarandGitHub 80a38aca2c feat(flags): Add local props and flags to all calls (#106) 2024-01-09 11:30:10 +00:00
Neil KakkarandGitHub 5159528070 feat(flags): Add relative date operators, fix numeric ops (#105) 2024-01-09 11:22:18 +00:00
Xavier VelloandGitHub e72101c03b feat: increase message and batch sizes (#108) 2023-12-04 16:27:26 +01:00
Daniil OkhlopkovandGitHub cf886de444 Update LICENSE (#104) 2023-10-25 16:08:07 +01:00
Neil KakkarandGitHub 8196da17cd fix(flags): Ensure feature properties exist on feature flag called ev… (#101)
* fix(flags): Ensure feature properties exist on feature flag called events

* fix

* make flake8 happy, will this bork?
2023-08-17 18:03:28 +01:00
Neil KakkarandGitHub e6116112f1 fix: disable behaviour for feature flags (#99) 2023-04-21 13:10:02 +01:00
fac1a3a16d Disable geoip for capture and decide calls by default (#98)
Co-authored-by: Neil Kakkar <neilkakkar@gmail.com>
2023-04-17 12:06:57 +01:00
Eric DuongandGitHub 2679ca6f38 version 2.5.0 (#97) 2023-04-10 15:13:25 -04:00
Eric DuongandGitHub 508182500b chore: change package to be an instantiable client (#96)
* initial try

* change init file to a class

* remove commas

* format

* sort

* restore original and add renamed client

* format

* move disabled

* format

* handle changing var for global instance

* add test
2023-04-10 14:33:07 -04:00
Neil KakkarandGitHub a543fb74f9 Fix analytics package dependencies (#95) 2023-03-30 10:34:59 -04:00
Neil KakkarandGitHub 9916f64e8e Add CI to ensure no prints (#94) 2023-03-17 15:22:59 +00:00
Neil KakkarandGitHub 22489ffe9f Release for print fix (#93) 2023-03-17 12:41:16 +00:00
Jann KleenandGitHub 289a26c2fe Remove print() call (#92) 2023-03-17 12:37:01 +00:00
Neil KakkarandGitHub ccee7b666b feat(flags): Enable local evaluation for all cohorts (#91) 2023-03-16 10:47:29 +00:00
Eric DuongandGitHub 5302613832 fix test (#90) 2023-02-14 14:24:10 -05:00
Eric DuongandGitHub a7804baa52 fix: add function for active variants (#89)
* add function for active variants

* format
2023-02-14 10:45:14 -05:00
Eric DuongandGitHub d226d469e3 add changelog entry and update versoin (#87) 2023-02-07 10:47:26 -05:00
Eric DuongandGitHub c078c0944c chore: Remove upper bound (#85)
* remove upper bound

* format
2023-02-07 10:36:26 -05:00
80523e03a2 fix: change api error to log (#83)
* change api error to log

* change implementation

* sort

* format

* raise when debug is true

* Update posthog/test/test_feature_flags.py

* format

---------

Co-authored-by: Neil Kakkar <neilkakkar@gmail.com>
2023-02-07 10:17:56 -05:00
Eric DuongandGitHub e76bfb16ee update version and changelog (#84) 2023-01-31 13:05:04 -05:00
Eric DuongandGitHub 57e30aca08 feat(feature-flags): JSON payload function (#81)
* implementation with test

* format

* v=3 and get_all_payloads

* format

* fix bug

* format

* address comments

* format

* add tests

* format

* add resiliency

* fix lookup

* remove check

* address comments

* example.py

* format

* format
2023-01-31 12:47:32 -05:00
Luke Harries 36b87e5d0b Revert "added regression tests"
This reverts commit 51f3cf6d06.
2023-01-24 17:31:30 +00:00
Luke Harries 51f3cf6d06 added regression tests 2023-01-24 17:11:37 +00:00
Neil KakkarandGitHub 34ee449a7e feat(flags): Add support for variant overrides (#77) 2022-11-14 12:15:33 +00:00
Neil KakkarandGitHub ce5aff1fe8 fix: datetime comparison issues (#75) 2022-09-15 15:26:38 +01:00
Neil KakkarandGitHub 93ae071693 Update CHANGELOG.md 2022-09-14 14:22:28 +01:00
Eric DuongandGitHub 5f1ba52ad5 feat: add date comparison to local evaluation (#74) 2022-09-14 13:53:44 +01:00
Neil KakkarandGitHub 5d6672315c fix: Remove defaults for feature flag calls (#72) 2022-08-12 12:20:54 +01:00
Neil KakkarandGitHub f490dd571a fix(feature-flags): Add more options to make using library easier (#70) 2022-08-04 13:53:20 +01:00
Neil KakkarandGitHub 78c7cb9ab2 Bump version to 2.0.0 and add breaking changes changelog (#69) 2022-08-02 12:10:32 +01:00
Neil KakkarandGitHub f0fd79841d feat(feature-flags): Enable local evaluation of flags (#68) 2022-07-29 14:10:51 +01:00
59dba76dea Add get feature flag method (#67)
* get feature flag method

* remove groups param from method

* use personal api key

* add groups

* add capture

* black reformat

* black?

* Update posthog/client.py

Co-authored-by: Neil Kakkar <neilkakkar@gmail.com>

* add to init

* formatting

Co-authored-by: Neil Kakkar <neilkakkar@gmail.com>
2022-06-30 11:48:53 -04:00
Neil KakkarandGitHub ccde499d53 bump version to 1.4.9 (#66) 2022-06-13 12:18:50 +01:00
62ea178d23 Have an option to send feature variants with the .capture(...) calls (#65)
Co-authored-by: Utku Zihnioglu <utku@webshare.io>
2022-06-13 12:13:35 +01:00
Tim GlaserandGitHub 0ec3a21498 bump version 1.4.8 (#61)
* bump version 1.4.8

* Update CHANGELOG.md
2022-05-12 08:12:11 +01:00
796e5df626 Feature: Enable multi variate feature flags for Python library (#60)
* Capturing $feature_flag_called at the end of client.feature_enabled method

* Enabling multi-variants for feature flags

Co-authored-by: Utku Zihnioglu <utku@webshare.io>
Co-authored-by: Tim Glaser <tim@posthog.com>
2022-05-12 08:10:54 +01:00
7f20286af7 Capturing $feature_flag_called at the end of client.feature_enabled method (#57)
Co-authored-by: Utku Zihnioglu <utku@webshare.io>
2022-05-12 08:05:26 +01:00
Tim Glaser 594b5f37e6 Bump version 1.4.7 2022-04-25 14:08:08 +01:00
Tim GlaserandGitHub a6ebb02d43 fix: Personal api key not required (#56)
* fix: Personal api key not required

* formatting

* Fix false
2022-04-25 14:07:08 +01:00
Tim GlaserandGitHub 9e31b3fdbb docs:Remove id: from example (#55) 2022-04-11 16:01:27 +02:00
Marius Andra 85feb84d87 1.4.6 2022-03-30 08:45:14 +02:00
Marius AndraandGitHub 6ca9b0002a Support custom UUID values (#53)
* feat(uuid): add support for custom uuids

* fix deprecated assert

* do not send "none" uuid

* you can have it in any color you want, as long as it's black
2022-03-30 08:42:43 +02:00
Paolo D'AmicoandGitHub 0723332b09 Update version.py (#49) 2022-01-05 06:46:28 -06:00
Paolo D'AmicoandGitHub c9ede4b953 Use Sentry DSN to obtain project ID (#48) 2022-01-05 06:15:51 -06:00
Karl-Aksel PuulmannandGitHub f7626714e2 Update version to 1.4.4 2021-11-25 11:45:18 +02:00
Karl-Aksel PuulmannandGitHub 8ac068d9a7 Groups: Feature flags support (#45)
* Fix a documentation typo

* Feature flags & groups support

* Update examples
2021-11-25 09:55:40 +02:00
Karl-Aksel PuulmannandGitHub 521a508c0b Update version to 1.4.3 2021-10-28 18:42:58 +03:00
Karl-Aksel PuulmannandGitHub 589475336b Basic group analytics support (#44)
* Make setup instructions work

* Add basic group analytics support to library

* Resolve formatting issues
2021-10-27 10:25:33 +03:00
Michael Matloka bd6f14a7cd Update pip install 2021-06-22 14:41:11 +02:00
Michael MatlokaandGitHub 41b3668b77 Update README.md 2021-06-22 14:39:02 +02:00
Michael MatlokaandGitHub c333f434fd Update README.md 2021-06-22 14:38:44 +02:00
Michael MatlokaandGitHub 4d3707d31c Bump version to 1.4.2 2021-06-22 14:38:15 +02:00
Michael Matloka d4fa6c6df7 Use PyPI API token instead of username and password 2021-06-22 14:30:53 +02:00
Michael Matloka 69d416e897 Ensure wheel 2021-06-22 14:24:27 +02:00
Michael Matloka e37c75e509 Fix make target in release workflow 2021-06-22 14:21:01 +02:00
Michael MatlokaandGitHub 0617277bd0 Add release workflow (#42)
* Add release workflow

* Update README.md
2021-06-22 14:19:31 +02:00
Yakko MajuriandGitHub 2b620e0f68 Fix integrity issues (#41)
* Fix integrity issues

* fix black, tests

* add test for new behavior

* run black properly

* remove accidental commit
2021-06-16 11:06:14 -03:00
Yakko MajuriandGitHub b239002483 Better handling on feature_enabled (#40) 2021-06-07 11:39:17 +01:00
Mandeep GillandGitHub e706401495 Remove pinning of backoff dependency (#38) 2021-06-04 14:26:36 -03:00
Neil Kakkar 7596674578 Release 1.4.1 2021-05-28 10:19:56 +01:00
Mandeep GillandGitHub b37047ab4d Bump backoff dependency to 1.10.0 (#36) 2021-05-26 19:02:15 +02:00
Michael MatlokaandGitHub 9df7af697b Fix 1.4.0 date in CHANGELOG.md 2021-05-20 15:44:53 +02:00
Michael MatlokaandGitHub 16685b1875 Update CHANGELOG.md 2021-05-18 17:21:59 +02:00
Michael MatlokaandGitHub 953140e207 Add PyPI badge to README 2021-05-18 17:20:36 +02:00
Michael Matloka 280e16fb2e Release 1.4.0 2021-05-18 17:16:37 +02:00
Neil KakkarandGitHub 4050c060d7 Add support for project_api_key. (#32)
* add support for project_api_key

* rm keys

* lint
2021-05-18 17:09:07 +02:00
Neil KakkarandGitHub fa2e69d158 Resolve polling issues with feature flags (#29)
* resolve polling issues

* lint
2021-05-17 14:09:20 +01:00
4c40410f4c Add python sentry (and sentry & django) integrations (#13)
* Add python sentry (and sentry & django) integrations

* example for django-sentry application and PosthogIntegration code

Co-authored-by: Neil Kakkar <neilkakkar@gmail.com>
2021-05-17 14:08:46 +01:00
Neil KakkarandGitHub 5261de0d29 Fix feature flag issue with no % rollout (#30)
* fix feature flag issue with no % rollout

* black
2021-05-14 11:52:06 -03:00
Neil Kakkar ca2673ddc5 Release 1.3.1 2021-05-07 12:22:14 +01:00
Neil KakkarandGitHub 2262a155c6 revert test change (#28) 2021-05-07 12:39:50 +02:00
Neil KakkarandGitHub 223e12efdd Merge pull request #23 from PostHog/set_once
Add $set and $set_once support
2021-05-07 11:06:04 +01:00
Neil Kakkar 1089749f72 black 2021-05-07 10:57:53 +01:00
Neil Kakkar 69eb0696c0 Add set and set_once to simulator.py 2021-05-07 10:54:10 +01:00
Tim GlaserandGitHub a26908d48d Merge pull request #27 from gagantrivedi/fix/alias
fix: Add distinct id to $create_alias event
2021-05-07 11:35:00 +02:00
Gagan a6350c9362 chore: Reformat consumer with black 2021-05-07 14:52:25 +05:30
Gagan a702f9ed94 fix: Add distinct id to $create_alias event 2021-05-07 11:22:03 +05:30
Neil Kakkar 0813f26a1d black 2021-05-06 14:05:57 +01:00
Neil Kakkar 7ecdba1a5d resolve conflicts with master 2021-05-06 13:59:46 +01:00
Neil Kakkar 8ecd73336b remove API key 2021-05-06 13:53:18 +01:00
Neil Kakkar 6d80cadf3c add functionality+tests for set and set_once 2021-05-06 13:52:23 +01:00
Neil Kakkar 5354f9aec8 cleanup tests, add instructions to readme 2021-05-06 11:43:13 +01:00
Paolo D'AmicoandGitHub 55ff3d94af Merge pull request #26 from PostHog/distinct-uuid 2021-04-01 17:59:29 -07:00
Michael Matloka d06c99a00c Add UUID to ID_TYPES 2021-04-01 17:31:48 +02:00
Paolo D'Amico f914eb4254 black 2021-02-08 11:21:55 +01:00
Paolo D'Amico 8e16d82553 update tests 2021-02-08 11:18:55 +01:00
Paolo D'Amico 0341dbc7db add $set_once support 2021-02-08 11:16:27 +01:00
Michael Matloka 46f8519310 Bump version to 1.2.1 2021-02-05 16:27:51 +01:00
Michael MatlokaandGitHub 37c7360b2c Merge pull request #22 from PostHog/fix-possibly-null-percentage
Add condition for rollout-percentage
2021-01-28 15:01:30 +01:00
Eric 1901eafbd9 add condition 2021-01-25 19:57:21 -05:00
Yakko MajuriandGitHub 07dfe8704f Update README.md 2021-01-11 17:19:21 -03:00
Michael MatlokaandGitHub aa435f373f Merge pull request #21 from PostHog/code-style
Black and isort all the things
2021-01-05 12:25:57 +01:00
Michael Matloka fc9d579791 Add CI job for this 2021-01-05 12:21:39 +01:00
Michael Matloka 93636181db Normalize strings to double quotes 2021-01-05 12:09:44 +01:00
Michael Matloka 78658eb034 Fix weird concat 2021-01-05 12:07:05 +01:00
Michael Matloka 3dbfb2f826 Black and isort all the things! 2021-01-03 05:37:51 +01:00
Michael Matloka 8761d1745a Setup black, isort and pre-commit 2021-01-03 05:36:26 +01:00
Michael MatlokaandGitHub c9ae5f73ee Merge pull request #20 from Hungsiro506/patch-1
Fix post() returning None on non-200 response
2021-01-03 05:07:51 +01:00
Michael Matloka 6f6a25f063 Update .gitignore 2021-01-03 05:00:41 +01:00
Michael Matloka dce9af2fae Add typing to request.py 2021-01-03 04:58:55 +01:00
Hưng VũandGitHub 92c80f75bd Update request.py
-- Fix post method return None
2020-12-30 17:43:08 +07:00
Michael MatlokaandGitHub 960941f534 Merge pull request #18 from PostHog/better-errors
Add project_api_key, improve errors
2020-12-15 12:01:28 +01:00
Michael MatlokaandGitHub fe3907d5f7 Merge pull request #14 from PostHog/uuid-property
Allow passing in UUID as property
2020-12-15 11:21:31 +01:00
Michael Matloka 54b6e83947 Rerun CI 2020-12-15 11:14:18 +01:00
yakkomajuri e93e0f9db0 update generic error handler 2020-12-04 11:10:22 +00:00
yakkomajuri 39670b2f01 specify json return 2020-12-04 10:29:07 +00:00
yakkomajuri caa33cd1b9 add project_api_key to config 2020-12-04 09:58:43 +00:00
yakkomajuri 7098a2d29a Better API error handling 2020-12-04 09:34:23 +00:00
Paolo D'AmicoandGitHub aac318f81f Merge pull request #17 from PostHog/v1.1.3 2020-11-23 09:33:50 -06:00
Paolo D'Amico 6882dbffbc version bump 2020-11-23 10:08:04 -05:00
Marius AndraandGitHub adf10b3d86 Do not start Feature Flag polling if no API Key (#15)
* jetbrains idea .gitignore

* do not start the poller if no personal api key
2020-11-23 15:37:52 +01:00
Michael Matloka 37b462120f Allow passing in UUID as property 2020-11-03 17:14:00 +01:00
Tim GlaserandGitHub 268384611a Merge pull request #12 from PostHog/investigate
Send distinctID to the decide endpoint to determin if user should have features enabled
2020-09-30 14:34:43 +01:00
James Greenhill c4cab81f90 bump version to 1.1.2 2020-09-30 14:25:49 +01:00
James Greenhill 1f51afc1b4 better testing and org 2020-09-30 14:20:50 +01:00
James Greenhill 57bc6320e2 fix tests 2020-09-30 14:10:34 +01:00
James Greenhill 6fa49c3f22 Send distinctID to the decide endpoint to determin if user should have features enabled 2020-09-30 13:53:47 +01:00
James GreenhillandGitHub c3695c5f13 Merge pull request #11 from PostHog/none
Default to False if the feature flag key does not exist. None is confusing
2020-09-29 17:06:53 +01:00
James Greenhill e6b84635f9 type annotations 2020-09-29 17:02:39 +01:00
James Greenhill 450383943b Default to False if the feature flag key does not exist. None is confusing 2020-09-29 16:07:43 +01:00
Tim GlaserandGitHub 053f20b846 Merge pull request #10 from PostHog/warn-instead-of-error
Warn instead of error
2020-09-29 13:16:06 +02:00
Tim Glaser 03f0223697 Fix tests 2020-09-29 12:13:28 +01:00
Tim Glaser df5f0003cd Warn instead of error 2020-09-29 11:57:44 +01:00
Tim GlaserandGitHub 9fed4e34f5 Merge pull request #9 from PostHog/feature-flags
Feature flags
2020-09-17 10:17:42 +02:00
Tim Glaser e5a4c896d0 Add error handling and timeouts 2020-09-17 10:13:03 +02:00
Yakko Majuri 421de30e0a minor changes 2020-09-14 09:37:14 +00:00
Yakko Majuri 84c05f2266 fix comment 2020-09-14 08:26:40 +00:00
Tim GlaserandGitHub c52dfdd81e Merge pull request #8 from PostHog/add-tests
Fix tests and clean up integration
2020-09-11 15:39:51 +02:00
Tim Glaser b42214f804 Merge branch 'add-tests' of github.com:PostHog/posthog-python into add-tests 2020-09-11 15:37:31 +02:00
Tim Glaser cd66eacac6 Cleaned up simulator 2020-09-11 15:37:06 +02:00
Yakko Majuri 4f2bad3fd2 fix action 2020-09-11 08:58:14 +00:00
Yakko Majuri 467b04403a minor changes 2020-09-11 08:50:50 +00:00
Tim Glaser 3bf1b5666c typo 2020-09-10 16:52:31 +02:00
Tim Glaser fdbebf4147 Fix tests 2020-09-10 16:39:27 +02:00
Tim Glaser 0d4a1c9444 Add poller 2020-09-10 16:35:23 +02:00
Tim Glaser c8526e15da Make feature flags work 2020-09-10 16:30:54 +02:00
Yakko Majuri 02a19c3d44 further cleanup of action 2020-09-09 16:48:07 +00:00
Yakko Majuri a8b0a7e12c updated tests action 2020-09-09 16:44:41 +00:00
Tim Glaser 730a84ae5e fix test 2020-09-09 16:26:59 +02:00
Tim Glaser 28e76606a9 Fix tests 2020-09-09 16:22:40 +02:00
Tim Glaser e34cf44205 Move tests into correct folder 2020-09-09 16:16:54 +02:00
Tim Glaser 45ae2891ce Fix tests 2020-09-09 16:15:46 +02:00
Marius AndraandGitHub 16e4d7e433 Update README.md 2020-05-20 16:43:32 +02:00
Tim Glaser 58d9e63834 Release 1.0.11 2020-04-30 10:10:51 +01:00
Tim GlaserandGitHub ca684c5d2c Merge pull request #6 from PostHog/5-timestamp-correct-type
Closes #5, add correct timestamp types
2020-04-30 10:08:51 +01:00
Tim Glaser a9e427bac4 Closes #5, add correct timestamp types 2020-04-30 10:08:28 +01:00
Marius AndraandGitHub 9c6595f66d Merge pull request #3 from casio/patch-1
Fix typo in identify() docstring
2020-04-29 22:12:44 +02:00
Marius AndraandGitHub 8df4c44be3 Merge pull request #4 from PostHog/links-fix
Fixed links to docs
2020-04-29 22:12:05 +02:00
James HawkinsandGitHub 34de69d736 Fixed links to docs 2020-04-07 16:26:06 +01:00
Tim Glaser ae3edee81f Update instructions 2020-04-01 17:26:10 +01:00
Tim Glaser e23c5a481b Release 1.0.10 2020-04-01 17:25:23 +01:00
Carsten KrausandGitHub 341e6763d3 Fix typo in docstring
lil typo, fixed : )
2020-03-25 22:42:05 +01:00
Tim Glaser 6ea2efcd9c Merge branch 'master' of github.com:PostHog/posthog-python 2020-02-21 17:09:32 -08:00
Tim GlaserandGitHub a514ff8ff2 Merge pull request #1 from mariusandra/patch-1
Fix typo with "identify" code example
2020-02-18 18:19:25 -08:00
Marius AndraandGitHub 023c33a7de Fix typo with "identify" code example 2020-02-18 22:10:02 +01:00
122 changed files with 30310 additions and 6385 deletions
+6 -6
View File
@@ -1,11 +1,11 @@
# PostHog API Configuration
# Hanzo Insights API Configuration
# Copy this file to .env and update with your actual values
# Your project API key (found on the /setup page in PostHog)
POSTHOG_PROJECT_API_KEY=phc_your_project_api_key_here
# Your project API key (found on the setup page in Insights)
INSIGHTS_PROJECT_API_KEY=hi_your_project_api_key_here
# Your personal API key (for local evaluation and other advanced features)
POSTHOG_PERSONAL_API_KEY=phx_your_personal_api_key_here
INSIGHTS_PERSONAL_API_KEY=phx_your_personal_api_key_here
# PostHog host URL (remove this line if using posthog.com)
POSTHOG_HOST=http://localhost:8000
# Insights host URL (remove this line if using insights.hanzo.ai)
INSIGHTS_HOST=http://localhost:8000
+28 -1
View File
@@ -36,6 +36,10 @@ jobs:
run: |
ruff format --check .
- name: Lint with ruff
run: |
ruff check .
- name: Check types with mypy
run: |
mypy --no-site-packages --config-file mypy.ini . | mypy-baseline filter
@@ -45,7 +49,7 @@ jobs:
runs-on: ubuntu-latest
strategy:
matrix:
python-version: ['3.9', '3.10', '3.11', '3.12', '3.13']
python-version: ['3.10', '3.11', '3.12', '3.13', '3.14']
steps:
- uses: actions/checkout@85e6279cec87321a52edac9c87bce653a07cf6c2
@@ -72,6 +76,29 @@ jobs:
run: |
pytest --verbose --timeout=30
import-check:
name: Python ${{ matrix.python-version }} import check
runs-on: ubuntu-latest
strategy:
matrix:
python-version: ['3.10', '3.11', '3.12', '3.13', '3.14']
steps:
- uses: actions/checkout@85e6279cec87321a52edac9c87bce653a07cf6c2
with:
fetch-depth: 1
- name: Set up Python ${{ matrix.python-version }}
uses: actions/setup-python@8d9ed9ac5c53483de85588cdf95a591a75ab9f55
with:
python-version: ${{ matrix.python-version }}
- name: Install posthog
run: pip install .
- name: Check import produces no warnings
run: python -W error -c "import posthog"
django5-integration:
name: Django 5 integration tests
runs-on: ubuntu-latest
+45
View File
@@ -0,0 +1,45 @@
name: 'CodeQL Advanced'
on:
push:
branches: ['master']
pull_request:
branches: ['master']
schedule:
- cron: '32 13 * * 1'
jobs:
analyze:
name: Analyze (${{ matrix.language }})
runs-on: 'ubuntu-latest'
permissions:
security-events: write
# required to fetch internal or private CodeQL packs
packages: read
strategy:
fail-fast: false
matrix:
include:
- language: actions
build-mode: none
- language: python
build-mode: none
steps:
- name: Checkout repository
uses: actions/checkout@v6
- name: Initialize CodeQL
uses: github/codeql-action/init@5d4e8d1aca955e8d8589aabd499c5cae939e33c7 # v4.31.9
with:
languages: ${{ matrix.language }}
build-mode: ${{ matrix.build-mode }}
# Disable TRAP caching - it creates a new cache per commit SHA which
# is never reused, causing wasted cache space.
# See: https://github.com/github/codeql-action/issues/2030
trap-caching: false
- name: Perform CodeQL Analysis
uses: github/codeql-action/analyze@5d4e8d1aca955e8d8589aabd499c5cae939e33c7 # v4.31.9
with:
category: '/language:${{matrix.language}}'
+2 -1
View File
@@ -7,12 +7,13 @@ jobs:
docs-generation:
name: Generate references
runs-on: ubuntu-latest
permissions:
contents: write
steps:
- name: Checkout the repository
uses: actions/checkout@85e6279cec87321a52edac9c87bce653a07cf6c2
with:
fetch-depth: 0
token: ${{ secrets.POSTHOG_BOT_PAT }}
- name: Set up Python
uses: actions/setup-python@8d9ed9ac5c53483de85588cdf95a591a75ab9f55
+225 -33
View File
@@ -1,57 +1,249 @@
name: "Release"
on:
push:
branches:
- master
paths:
- "posthog/version.py"
pull_request:
types: [closed]
branches: [master]
workflow_dispatch:
permissions:
contents: read
# Concurrency control: only one release process can run at a time
# This prevents race conditions if multiple PRs with 'release' label merge simultaneously
concurrency:
group: release
cancel-in-progress: false
jobs:
release:
name: Publish release
check-release-label:
name: Check for release label
runs-on: ubuntu-latest
env:
TWINE_USERNAME: __token__
TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }}
# Run when PR with 'release' label is merged to master
if: |
github.event_name == 'workflow_dispatch' ||
(github.event_name == 'pull_request' &&
github.event.pull_request.merged == true &&
contains(github.event.pull_request.labels.*.name, 'release'))
outputs:
should-release: ${{ steps.check.outputs.should-release }}
steps:
- name: Checkout the repository
uses: actions/checkout@85e6279cec87321a52edac9c87bce653a07cf6c2
- name: Checkout repository
uses: actions/checkout@v4
with:
ref: master
fetch-depth: 0
token: ${{ secrets.POSTHOG_BOT_PAT }}
- name: Check release conditions
id: check
run: |
changeset_count=$(find .sampo/changesets -name '*.md' 2>/dev/null | wc -l)
if [ "$changeset_count" -gt 0 ]; then
echo "should-release=true" >> "$GITHUB_OUTPUT"
echo "Found $changeset_count changeset(s), ready to release"
else
echo "should-release=false" >> "$GITHUB_OUTPUT"
echo "No changesets to release"
fi
notify-approval-needed:
name: Notify Slack - Approval Needed
needs: check-release-label
if: needs.check-release-label.outputs.should-release == 'true'
uses: posthog/.github/.github/workflows/notify-approval-needed.yml@main
with:
slack_channel_id: ${{ vars.SLACK_APPROVALS_CLIENT_LIBRARIES_CHANNEL_ID }}
slack_user_group_id: ${{ vars.GROUP_CLIENT_LIBRARIES_SLACK_GROUP_ID }}
secrets:
slack_bot_token: ${{ secrets.SLACK_CLIENT_LIBRARIES_BOT_TOKEN }}
posthog_project_api_key: ${{ secrets.POSTHOG_PROJECT_API_KEY }}
release:
name: Release and publish
needs: [check-release-label, notify-approval-needed]
runs-on: ubuntu-latest
# Use `always()` to ensure the job runs even if notify-approval-needed is skipped,
# but still depend on it to access `needs.notify-approval-needed.outputs.slack_ts`
if: always() && needs.check-release-label.outputs.should-release == 'true'
environment: "Release" # This will require an approval from a maintainer, they are notified in Slack above
permissions:
contents: write
actions: write
id-token: write
steps:
- name: Notify Slack - Approved
if: needs.notify-approval-needed.outputs.slack_ts != ''
uses: posthog/.github/.github/actions/slack-thread-reply@main
with:
slack_bot_token: ${{ secrets.SLACK_CLIENT_LIBRARIES_BOT_TOKEN }}
slack_channel_id: ${{ vars.SLACK_APPROVALS_CLIENT_LIBRARIES_CHANNEL_ID }}
thread_ts: ${{ needs.notify-approval-needed.outputs.slack_ts }}
message: "✅ Release approved! Version bump in progress..."
emoji_reaction: "white_check_mark"
- name: Get GitHub App token
id: releaser
uses: actions/create-github-app-token@v2
with:
app-id: ${{ secrets.GH_APP_POSTHOG_PYTHON_RELEASER_APP_ID }}
private-key: ${{ secrets.GH_APP_POSTHOG_PYTHON_RELEASER_PRIVATE_KEY }}
- name: Checkout repository
uses: actions/checkout@v4
with:
ref: master
fetch-depth: 0
token: ${{ steps.releaser.outputs.token }}
- name: Set up Python
uses: actions/setup-python@8d9ed9ac5c53483de85588cdf95a591a75ab9f55
uses: actions/setup-python@v5
with:
python-version: 3.11.11
- name: Install uv
uses: astral-sh/setup-uv@0c5e2b8115b80b4c7c5ddf6ffdd634974642d182 # v5.4.1
uses: astral-sh/setup-uv@v5
with:
enable-cache: true
pyproject-file: 'pyproject.toml'
- name: Detect version
run: echo "REPO_VERSION=$(python3 posthog/version.py)" >> $GITHUB_ENV
enable-cache: true
pyproject-file: "pyproject.toml"
- name: Prepare for building release
- name: Install Rust
uses: dtolnay/rust-toolchain@0b1efabc08b657293548b77fb76cc02d26091c7e
with:
toolchain: 1.91.1
components: cargo
- name: Cache Sampo CLI
id: cache-sampo
uses: actions/cache@v3
with:
path: ~/.cargo/bin/sampo
key: sampo-${{ runner.os }}-${{ runner.arch }}
- name: Install Sampo CLI
if: steps.cache-sampo.outputs.cache-hit != 'true'
run: cargo install sampo
- name: Install dependencies
run: uv sync --extra dev
- name: Push releases to PyPI
run: uv run make release && uv run make release_analytics
- name: Configure Git
run: |
git config user.name "github-actions[bot]"
git config user.email "github-actions[bot]@users.noreply.github.com"
- name: Create GitHub release
uses: actions/create-release@0cb9c9b65d5d1901c1f53e5e66eaf4afd303e70e # v1
- name: Prepare release with Sampo
id: sampo-release
env:
GITHUB_TOKEN: ${{ secrets.POSTHOG_BOT_PAT }}
with:
tag_name: v${{ env.REPO_VERSION }}
release_name: ${{ env.REPO_VERSION }}
- name: Dispatch generate-references for posthog-python
GITHUB_TOKEN: ${{ steps.releaser.outputs.token }}
run: |
sampo release
new_version=$(python3 -c "import tomllib; print(tomllib.load(open('pyproject.toml', 'rb'))['project']['version'])")
echo "new_version=$new_version" >> "$GITHUB_OUTPUT"
- name: Sync version to posthog/version.py
run: |
echo 'VERSION = "${{ steps.sampo-release.outputs.new_version }}"' > posthog/version.py
- name: Commit release changes
id: commit-release
env:
GITHUB_TOKEN: ${{ steps.releaser.outputs.token }}
run: |
git add -A
if git diff --staged --quiet; then
echo "No changes to commit"
echo "committed=false" >> "$GITHUB_OUTPUT"
else
git commit -m "chore: Release v${{ steps.sampo-release.outputs.new_version }}"
git push origin master
echo "committed=true" >> "$GITHUB_OUTPUT"
fi
# Publishing is done manually (not via `sampo publish`) because we need to
# publish both `posthog` and `posthoganalytics` packages to PyPI.
# Sampo only knows about the `posthog` package, so we handle both here.
# Both packages use PyPI OIDC trusted publishing (no API tokens needed).
- name: Build posthog
if: steps.commit-release.outputs.committed == 'true'
run: uv run make build_release
- name: Publish posthog to PyPI
if: steps.commit-release.outputs.committed == 'true'
uses: pypa/gh-action-pypi-publish@release/v1
# The `posthoganalytics` package is a mirror of `posthog` published under
# a different name for backwards compatibility. The make target handles
# copying, renaming imports, and building the dist automatically.
- name: Build posthoganalytics
if: steps.commit-release.outputs.committed == 'true'
run: uv run make build_release_analytics
- name: Publish posthoganalytics to PyPI
if: steps.commit-release.outputs.committed == 'true'
uses: pypa/gh-action-pypi-publish@release/v1
# We skip `sampo publish` (which normally creates the tag) because we
# need to publish both posthog and posthoganalytics manually, so we
# create the tag ourselves.
- name: Tag release
if: steps.commit-release.outputs.committed == 'true'
run: git tag "v${{ steps.sampo-release.outputs.new_version }}"
- name: Push tags
if: steps.commit-release.outputs.committed == 'true'
run: git push origin --tags
- name: Create GitHub Release
if: steps.commit-release.outputs.committed == 'true'
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
run: gh release create "v${{ steps.sampo-release.outputs.new_version }}" --generate-notes
- name: Dispatch generate-references
if: steps.commit-release.outputs.committed == 'true'
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
run: |
gh workflow run generate-references.yml --ref master
run: gh workflow run generate-references.yml --ref master
# Notify in case of a failure
- name: Send failure event to PostHog
if: ${{ failure() }}
uses: PostHog/posthog-github-action@v0.1
with:
posthog-token: "${{ secrets.POSTHOG_PROJECT_API_KEY }}"
event: "posthog-python-github-release-workflow-failure"
properties: >-
{
"commitSha": "${{ github.sha }}",
"jobStatus": "${{ job.status }}",
"ref": "${{ github.ref }}",
"version": "v${{ steps.sampo-release.outputs.new_version }}"
}
- name: Notify Slack - Failed
if: ${{ failure() && needs.notify-approval-needed.outputs.slack_ts != '' }}
uses: posthog/.github/.github/actions/slack-thread-reply@main
with:
slack_bot_token: ${{ secrets.SLACK_CLIENT_LIBRARIES_BOT_TOKEN }}
slack_channel_id: ${{ vars.SLACK_APPROVALS_CLIENT_LIBRARIES_CHANNEL_ID }}
thread_ts: ${{ needs.notify-approval-needed.outputs.slack_ts }}
message: "❌ Failed to release `posthog-python@v${{ steps.sampo-release.outputs.new_version }}`! <https://github.com/${{ github.repository }}/actions/runs/${{ github.run_id }}|View logs>"
emoji_reaction: "x"
notify-released:
name: Notify Slack - Released
needs: [check-release-label, notify-approval-needed, release]
runs-on: ubuntu-latest
if: always() && needs.release.result == 'success' && needs.notify-approval-needed.outputs.slack_ts != ''
steps:
- name: Checkout repository
uses: actions/checkout@v4
- name: Notify Slack - Released
uses: posthog/.github/.github/actions/slack-thread-reply@main
with:
slack_bot_token: ${{ secrets.SLACK_CLIENT_LIBRARIES_BOT_TOKEN }}
slack_channel_id: ${{ vars.SLACK_APPROVALS_CLIENT_LIBRARIES_CHANNEL_ID }}
thread_ts: ${{ needs.notify-approval-needed.outputs.slack_ts }}
message: "🚀 posthog-python released successfully!"
emoji_reaction: "rocket"
+21
View File
@@ -0,0 +1,21 @@
name: SDK Compliance Tests
permissions:
contents: read
packages: read
pull-requests: write
on:
pull_request:
push:
branches:
- master
jobs:
compliance:
name: PostHog SDK compliance tests
uses: PostHog/posthog-sdk-test-harness/.github/workflows/test-sdk-action.yml@main
with:
adapter-dockerfile: "sdk_compliance_adapter/Dockerfile"
adapter-context: "."
test-harness-version: "latest"
+5
View File
@@ -0,0 +1,5 @@
---
pypi/posthog: patch
---
feat(llma): support fetching versioned prompts from the prompts sdk
+19
View File
@@ -0,0 +1,19 @@
# Sampo configuration
version = 1
[git]
default_branch = "master"
short_tags = "posthog" # Tag with v1.2.3 rather than posthog-v1.2.3
[github]
repository = "posthog/posthog-python"
[changelog]
# Options for release notes generation.
# show_commit_hash = true (default)
# show_acknowledgments = true (default)
[packages]
# Options for package discovery and filtering.
# ignore_unpublished = false (default)
# ignore = ["internal-*", "examples/*"]
Symlink
+1
View File
@@ -0,0 +1 @@
LLM.md
+7 -7
View File
@@ -1,6 +1,6 @@
# Before Send Hook
The `before_send` parameter allows you to modify or filter events before they are sent to PostHog. This is useful for:
The `before_send` parameter allows you to modify or filter events before they are sent to Insights. This is useful for:
- **Privacy**: Removing or masking sensitive data (PII)
- **Filtering**: Dropping unwanted events (test events, internal users, etc.)
@@ -10,12 +10,12 @@ The `before_send` parameter allows you to modify or filter events before they ar
## Basic Usage
```python
import posthog
import hanzo_insights
from typing import Optional, Dict, Any
def my_before_send(event: Dict[str, Any]) -> Optional[Dict[str, Any]]:
"""
Process event before sending to PostHog.
Process event before sending to Insights.
Args:
event: The event dictionary containing 'event', 'distinct_id', 'properties', etc.
@@ -27,7 +27,7 @@ def my_before_send(event: Dict[str, Any]) -> Optional[Dict[str, Any]]:
return event
# Initialize client with before_send hook
client = posthog.Client(
client = hanzo_insights.Client(
api_key="your-project-api-key",
before_send=my_before_send
)
@@ -166,7 +166,7 @@ def should_drop_event(event: dict[str, Any]) -> bool:
## Error Handling
If your `before_send` function raises an exception, PostHog will:
If your `before_send` function raises an exception, Insights will:
1. Log the error
2. Continue with the original, unmodified event
@@ -184,7 +184,7 @@ def risky_before_send(event: dict[str, Any]) -> Optional[dict[str, Any]]:
## Complete Example
```python
import posthog
import hanzo_insights
from typing import Optional, Any
import re
@@ -227,7 +227,7 @@ def production_before_send(event: dict[str, Any]) -> Optional[dict[str, Any]]:
return event # Return original event on error
# Usage
client = posthog.Client(
client = hanzo_insights.Client(
api_key="your-api-key",
before_send=production_before_send
)
+209 -38
View File
@@ -1,144 +1,315 @@
# 6.9.1 - 2025-11-07
# posthog
## 7.9.7 — 2026-03-05
### Patch changes
- [b206669](https://github.com/posthog/posthog-python/commit/b206669bf62c923346ad28881dc4694d933ca424) fix(llma): use distinct_id from outer context if not provided, fix $process_person_profile for context-based identity — Thanks @ethanporcaro for your first contribution 🎉!
- [a99c7d7](https://github.com/posthog/posthog-python/commit/a99c7d73b1e0ef1f35d856c82ace21237ee253a3) Add warning log for local flag evaluation cold start — Thanks @dmarticus!
## 7.9.6 — 2026-03-02
### Patch changes
- [8d83315](https://github.com/posthog/posthog-python/commit/8d83315b67c21eb9e7d6c17bae27ada98ca2643d) add PROPERTY_OPERATORS constant for match_property — Thanks @dmarticus!
## 7.9.5 — 2026-03-02
### Patch changes
- [830244b](https://github.com/posthog/posthog-python/commit/830244bd409b1992ae2e49610f8f87d2cdfc8096) add semver targeting support to local evaluation — Thanks @dmarticus!
## 7.9.4 — 2026-02-25
### Patch changes
- [a68a6a6](https://github.com/posthog/posthog-python/commit/a68a6a6d045072c88eeee7acac441536919b5954) feat(llma): add `$ai_tokens_source` property ("sdk" or "passthrough") to all `$ai_generation` events to detect when token values are externally overridden via `posthog_properties` — Thanks @carlos-marchal-ph!
## 7.9.3 — 2026-02-18
### Patch changes
- [9f9553a](https://github.com/posthog/posthog-python/commit/9f9553a420d22e5e6435b775993f61a059280c2a) Fix posthoganalytics release, previously broken — Thanks @rafaeelaudibert!
## 7.9.2 — 2026-02-18
### Patch changes
- [f1dc4d7](https://github.com/posthog/posthog-python/commit/f1dc4d73914712983a7f715ee4fe1b70e66e770a) Add sampo to the project — Thanks @rafaeelaudibert!
## 7.9.1 - 2026-02-17
fix(llma): make prompt fetches deterministic by requiring project_api_key and sending it as token query param
## 7.9.0 - 2026-02-17
feat: Support device_id as bucketing identifier for local evaluation
## 7.8.6 - 2026-02-09
fix: limit collections scanning in code variables
## 7.8.5 - 2026-02-09
fix: further optimize code variables pattern matching
## 7.8.4 - 2026-02-09
fix: do not pattern match long values in code variables
## 7.8.3 - 2026-02-06
fix: openAI input image sanitization
## 7.8.2 - 2026-02-04
fix(llma): fix prompts default url
## 7.8.1 - 2026-02-03
fix(llma): small fixes for prompt management
## 7.8.0 - 2026-01-28
feat(llma): add prompt management
Adds the Prompt Management feature. At the time of release, this feature is in a closed alpha.
## 7.7.0 - 2026-01-15
feat(ai): Add OpenAI Agents SDK integration
Automatic tracing for agent workflows, handoffs, tool calls, guardrails, and custom spans. Includes `$ai_total_tokens`, `$ai_error_type` categorization, and `$ai_framework` property.
## 7.6.0 - 2026-01-12
feat: add device_id to flags request payload
Add device_id parameter to all feature flag methods, allowing the server to track device identifiers for flag evaluation. The device_id can be passed explicitly or set via context using `set_context_device_id()`.
## 7.5.1 - 2026-01-07
fix: avoid return from finally block to fix Python 3.14 SyntaxWarning (#361) - thanks @jodal
## 7.5.0 - 2026-01-06
feat: Capture Langchain, OpenAI and Anthropic errors as exceptions (if exception autocapture is enabled)
feat: Add reference to exception in LLMA trace and span events
## 7.4.3 - 2026-01-02
Fixes cache creation cost for Langchain with Anthropic
## 7.4.2 - 2025-12-22
feat: add `in_app_modules` option to control code variables capturing
## 7.4.1 - 2025-12-19
fix: extract model from response for OpenAI stored prompts
When using OpenAI stored prompts, the model is defined in the OpenAI dashboard rather than passed in the API request. This fix adds a fallback to extract the model from the response object when not provided in kwargs, ensuring generations show up with the correct model and enabling cost calculations.
## 7.4.0 - 2025-12-16
feat: Add automatic retries for feature flag requests
Feature flag API requests now automatically retry on transient failures:
- Network errors (connection refused, DNS failures, timeouts)
- Server errors (500, 502, 503, 504)
- Up to 2 retries with exponential backoff (0.5s, 1s delays)
Rate limit (429) and quota (402) errors are not retried.
## 7.3.1 - 2025-12-06
fix: remove unused $exception_message and $exception_type
## 7.3.0 - 2025-12-05
feat: improve code variables capture masking
## 7.2.0 - 2025-12-01
feat: add $feature_flag_evaluated_at properties to $feature_flag_called events
## 7.1.0 - 2025-11-26
Add support for the async version of Gemini.
## 7.0.2 - 2025-11-18
Add support for Python 3.14.
Projects upgrading to Python 3.14 should ensure any Pydantic models passed into the SDK use Pydantic v2, as Pydantic v1 is not compatible with Python 3.14.
## 7.0.1 - 2025-11-15
Try to use repr() when formatting code variables
## 7.0.0 - 2025-11-11
NB Python 3.9 is no longer supported
- chore(llma): update LLM provider SDKs to latest major versions
- openai: 1.102.0 → 2.7.1
- anthropic: 0.64.0 → 0.72.0
- google-genai: 1.32.0 → 1.49.0
- langchain-core: 0.3.75 → 1.0.3
- langchain-openai: 0.3.32 → 1.0.2
- langchain-anthropic: 0.3.19 → 1.0.1
- langchain-community: 0.3.29 → 0.4.1
- langgraph: 0.6.6 → 1.0.2
## 6.9.3 - 2025-11-10
- feat(ph-ai): PostHog properties dict in GenerationMetadata
## 6.9.2 - 2025-11-10
- fix(llma): fix cache token double subtraction in Langchain for non-Anthropic providers causing negative costs
## 6.9.1 - 2025-11-07
- fix(error-tracking): pass code variables config from init to client
# 6.9.0 - 2025-11-06
## 6.9.0 - 2025-11-06
- feat(error-tracking): add local variables capture
# 6.8.0 - 2025-11-03
## 6.8.0 - 2025-11-03
- feat(llma): send web search calls to be used for LLM cost calculations
# 6.7.14 - 2025-11-03
## 6.7.14 - 2025-11-03
- fix(django): Handle request.user access in async middleware context to prevent SynchronousOnlyOperation errors in Django 5+ (fixes #355)
- test(django): Add Django 5 integration test suite with real ASGI application testing async middleware behavior
# 6.7.13 - 2025-11-02
## 6.7.13 - 2025-11-02
- fix(llma): cache cost calculation in the LangChain callback
# 6.7.12 - 2025-11-02
## 6.7.12 - 2025-11-02
- fix(django): Restore process_exception method to capture view and downstream middleware exceptions (fixes #329)
- fix(ai/langchain): Add LangChain 1.0+ compatibility for CallbackHandler imports (fixes #362)
# 6.7.11 - 2025-10-28
## 6.7.11 - 2025-10-28
- feat(ai): Add `$ai_framework` property for framework integrations (e.g. LangChain)
# 6.7.10 - 2025-10-24
## 6.7.10 - 2025-10-24
- fix(django): Make middleware truly hybrid - compatible with both sync (WSGI) and async (ASGI) Django stacks without breaking sync-only deployments
# 6.7.9 - 2025-10-22
## 6.7.9 - 2025-10-22
- fix(flags): multi-condition flags with static cohorts returning wrong variants
# 6.7.8 - 2025-10-16
## 6.7.8 - 2025-10-16
- fix(llma): missing async for OpenAI's streaming implementation
# 6.7.7 - 2025-10-14
## 6.7.7 - 2025-10-14
- fix: remove deprecated attribute $exception_personURL from exception events
# 6.7.6 - 2025-09-16
## 6.7.6 - 2025-09-16
- fix: don't sort condition sets with variant overrides to the top
- fix: Prevent core Client methods from raising exceptions
# 6.7.5 - 2025-09-16
## 6.7.5 - 2025-09-16
- feat: Django middleware now supports async request handling.
# 6.7.4 - 2025-09-05
## 6.7.4 - 2025-09-05
- fix: Missing system prompts for some providers
# 6.7.3 - 2025-09-04
## 6.7.3 - 2025-09-04
- fix: missing usage tokens in Gemini
# 6.7.2 - 2025-09-03
## 6.7.2 - 2025-09-03
- fix: tool call results in streaming providers
# 6.7.1 - 2025-09-01
## 6.7.1 - 2025-09-01
- fix: Add base64 inline image sanitization
# 6.7.0 - 2025-08-26
## 6.7.0 - 2025-08-26
- feat: Add support for feature flag dependencies
# 6.6.1 - 2025-08-21
## 6.6.1 - 2025-08-21
- fix: Prevent `NoneType` error when `group_properties` is `None`
# 6.6.0 - 2025-08-15
## 6.6.0 - 2025-08-15
- feat: Add `flag_keys_to_evaluate` parameter to optimize feature flag evaluation performance by only evaluating specified flags
- feat: Add `flag_keys_filter` option to `send_feature_flags` for selective flag evaluation in capture events
# 6.5.0 - 2025-08-08
## 6.5.0 - 2025-08-08
- feat: Add `$context_tags` to an event to know which properties were included as tags
# 6.4.1 - 2025-08-06
## 6.4.1 - 2025-08-06
- fix: Always pass project API key in `remote_config` requests for deterministic project routing
# 6.4.0 - 2025-08-05
## 6.4.0 - 2025-08-05
- feat: support Vertex AI for Gemini
# 6.3.4 - 2025-08-04
## 6.3.4 - 2025-08-04
- fix: set `$ai_tools` for all providers and `$ai_output_choices` for all non-streaming provider flows properly
# 6.3.3 - 2025-08-01
## 6.3.3 - 2025-08-01
- fix: `get_feature_flag_result` now correctly returns FeatureFlagResult when payload is empty string instead of None
# 6.3.2 - 2025-07-31
## 6.3.2 - 2025-07-31
- fix: Anthropic's tool calls are now handled properly
# 6.3.0 - 2025-07-22
## 6.3.0 - 2025-07-22
- feat: Enhanced `send_feature_flags` parameter to accept `SendFeatureFlagsOptions` object for declarative control over local/remote evaluation and custom properties
# 6.2.1 - 2025-07-21
## 6.2.1 - 2025-07-21
- feat: make `posthog_client` an optional argument in PostHog AI providers wrappers (`posthog.ai.*`), intuitively using the default client as the default
# 6.1.1 - 2025-07-16
## 6.1.1 - 2025-07-16
- fix: correctly capture exceptions processed by Django from views or middleware
# 6.1.0 - 2025-07-10
## 6.1.0 - 2025-07-10
- feat: decouple feature flag local evaluation from personal API keys; support decrypting remote config payloads without relying on the feature flags poller
# 6.0.4 - 2025-07-09
## 6.0.4 - 2025-07-09
- fix: add POSTHOG_MW_CLIENT setting to django middleware, to support custom clients for exception capture.
# 6.0.3 - 2025-07-07
## 6.0.3 - 2025-07-07
- feat: add a feature flag evaluation cache (local storage or redis) to support returning flag evaluations when the service is down
# 6.0.2 - 2025-07-02
## 6.0.2 - 2025-07-02
- fix: send_feature_flags changed to default to false in `Client::capture_exception`
# 6.0.1
## 6.0.1
- fix: response `$process_person_profile` property when passed to capture
# 6.0.0
## 6.0.0
This release contains a number of major breaking changes:
@@ -165,15 +336,15 @@ with posthog.new_context():
Generally, arguments are now appropriately typed, and docstrings have been updated. If something is unclear, please open an issue, or submit a PR!
# 5.4.0 - 2025-06-20
## 5.4.0 - 2025-06-20
- feat: add support to session_id context on page method
# 5.3.0 - 2025-06-19
## 5.3.0 - 2025-06-19
- fix: safely handle exception values
# 5.2.0 - 2025-06-19
## 5.2.0 - 2025-06-19
- feat: construct artificial stack traces if no traceback is available on a captured exception
Symlink
+1
View File
@@ -0,0 +1 @@
LLM.md
+48
View File
@@ -0,0 +1,48 @@
# LLM.md - Hanzo Insights Python SDK
## Overview
Integrate Hanzo Insights into any Python application. Package name: `hanzo-insights` on PyPI.
## Tech Stack
- **Language**: Python 3.10+
- **Package**: `hanzo_insights` (import name), `hanzo-insights` (pip name)
## Build & Run
```bash
uv sync
uv run pytest
```
## Structure
```
posthog-python/
hanzo_insights/ # Main package
__init__.py # Module-level API, Insights class
client.py # Client class
ai/ # AI provider integrations (OpenAI, Anthropic, Gemini, LangChain)
integrations/ # Framework integrations (Django middleware)
test/ # Tests
examples/
integration_tests/
pyproject.toml # Package config (name: hanzo-insights)
setup.py # Legacy setup
```
## Key Files
- `pyproject.toml` -- Package config, dependencies, test config
- `hanzo_insights/__init__.py` -- Public API surface
- `hanzo_insights/client.py` -- Client implementation
## Rebrand Notes
- Main class: `Insights` (no backward compat aliases)
- Django middleware: `InsightsContextMiddleware` (no backward compat aliases)
- OpenAI Agents: `InsightsTracingProcessor` (no backward compat aliases)
- `$lib` protocol value: `insights-python`
- Ingestion URLs: `us.i.insights.hanzo.ai` / `eu.i.insights.hanzo.ai`
- AI wrapper kwargs: `insights_*` (e.g. `insights_distinct_id`, `insights_trace_id`)
- Exception attrs: `__insights_exception_captured`, `__insights_exception_uuid`
- Context var: `insights_context_stack`
- Redis prefix: `insights:flags:`
- Redaction sentinels: `$$_insights_redacted_*`, `$$_insights_value_too_long_*`
- Django settings: `INSIGHTS_MW_*` only (no `POSTHOG_MW_*` fallback)
- Django headers: `X-INSIGHTS-SESSION-ID`, `X-INSIGHTS-DISTINCT-ID` only
+32 -18
View File
@@ -5,28 +5,42 @@ test:
coverage run -m pytest
coverage report
release:
build_release:
rm -rf dist/*
python setup.py sdist bdist_wheel
twine upload dist/*
release_analytics:
# Builds the `posthoganalytics` PyPI package, which is a mirror of `hanzo_insights`
# published under a different name for backward compatibility with the upstream
# posthog/posthog project.
#
# The process works in three phases:
# 1. hanzo_insights -> posthoganalytics: Copy the source, rewrite all imports,
# remove the original hanzo_insights/ dir, and build the dist.
# 2. posthoganalytics -> hanzo_insights: Reverse the import rewrites, copy
# everything back into hanzo_insights/, and clean up.
# 3. Restore pyproject.toml from backup (setup_analytics.py modifies it).
#
# This ensures the working tree is left in the same state it started in.
#
# NOTE: This target clears dist/ before building. In the release workflow,
# `build_release` (hanzo_insights) must be published BEFORE running this target,
# otherwise the hanzo_insights dist artifacts will be lost.
build_release_analytics:
rm -rf dist
rm -rf build
rm -rf posthoganalytics
mkdir posthoganalytics
cp -r posthog/* posthoganalytics/
find ./posthoganalytics -type f -name "*.py" -exec sed -i.bak -e 's/from posthog /from posthoganalytics /g' {} \;
find ./posthoganalytics -type f -name "*.py" -exec sed -i.bak -e 's/from posthog\./from posthoganalytics\./g' {} \;
cp -r hanzo_insights/* posthoganalytics/
find ./posthoganalytics -type f -name "*.py" -exec sed -i.bak -e 's/from hanzo_insights /from posthoganalytics /g' {} \;
find ./posthoganalytics -type f -name "*.py" -exec sed -i.bak -e 's/from hanzo_insights\./from posthoganalytics\./g' {} \;
find ./posthoganalytics -name "*.bak" -delete
rm -rf posthog
rm -rf hanzo_insights
python setup_analytics.py sdist bdist_wheel
twine upload dist/*
mkdir posthog
find ./posthoganalytics -type f -name "*.py" -exec sed -i.bak -e 's/from posthoganalytics /from posthog /g' {} \;
find ./posthoganalytics -type f -name "*.py" -exec sed -i.bak -e 's/from posthoganalytics\./from posthog\./g' {} \;
mkdir hanzo_insights
find ./posthoganalytics -type f -name "*.py" -exec sed -i.bak -e 's/from posthoganalytics /from hanzo_insights /g' {} \;
find ./posthoganalytics -type f -name "*.py" -exec sed -i.bak -e 's/from posthoganalytics\./from hanzo_insights\./g' {} \;
find ./posthoganalytics -name "*.bak" -delete
cp -r posthoganalytics/* posthog/
cp -r posthoganalytics/* hanzo_insights/
rm -rf posthoganalytics
rm -f pyproject.toml
cp pyproject.toml.backup pyproject.toml
@@ -41,17 +55,17 @@ prep_local:
cp -r . ../posthog-python-local/
cd ../posthog-python-local && rm -rf dist build posthoganalytics .git
cd ../posthog-python-local && mkdir posthoganalytics
cd ../posthog-python-local && cp -r posthog/* posthoganalytics/
cd ../posthog-python-local && find ./posthoganalytics -type f -name "*.py" -exec sed -i.bak -e 's/from posthog /from posthoganalytics /g' {} \;
cd ../posthog-python-local && find ./posthoganalytics -type f -name "*.py" -exec sed -i.bak -e 's/from posthog\./from posthoganalytics\./g' {} \;
cd ../posthog-python-local && cp -r hanzo_insights/* posthoganalytics/
cd ../posthog-python-local && find ./posthoganalytics -type f -name "*.py" -exec sed -i.bak -e 's/from hanzo_insights /from posthoganalytics /g' {} \;
cd ../posthog-python-local && find ./posthoganalytics -type f -name "*.py" -exec sed -i.bak -e 's/from hanzo_insights\./from posthoganalytics\./g' {} \;
cd ../posthog-python-local && find ./posthoganalytics -name "*.bak" -delete
cd ../posthog-python-local && rm -rf posthog
cd ../posthog-python-local && rm -rf hanzo_insights
cd ../posthog-python-local && sed -i.bak 's/from version import VERSION/from posthoganalytics.version import VERSION/' setup_analytics.py
cd ../posthog-python-local && rm setup_analytics.py.bak
cd ../posthog-python-local && sed -i.bak 's/"posthog"/"posthoganalytics"/' setup.py
cd ../posthog-python-local && sed -i.bak 's/"hanzo_insights"/"posthoganalytics"/' setup.py
cd ../posthog-python-local && rm setup.py.bak
cd ../posthog-python-local && python -c "import setup_analytics" 2>/dev/null || true
@echo "Local copy created at ../posthog-python-local"
@echo "Install with: pip install -e ../posthog-python-local"
.PHONY: test lint release e2e_test prep_local
.PHONY: test lint build_release build_release_analytics e2e_test prep_local
+54 -50
View File
@@ -1,37 +1,55 @@
# PostHog Python
# Hanzo Insights Python SDK
<p align="center">
<img alt="posthoglogo" src="https://user-images.githubusercontent.com/65415371/205059737-c8a4f836-4889-4654-902e-f302b187b6a0.png">
</p>
<p align="center">
<a href="https://pypi.org/project/posthog/"><img alt="pypi installs" src="https://img.shields.io/pypi/v/posthog"/></a>
<img alt="GitHub contributors" src="https://img.shields.io/github/contributors/posthog/posthog-python">
<img alt="GitHub commit activity" src="https://img.shields.io/github/commit-activity/m/posthog/posthog-python"/>
<img alt="GitHub closed issues" src="https://img.shields.io/github/issues-closed/posthog/posthog-python"/>
</p>
Integrate [Hanzo Insights](https://insights.hanzo.ai) into any Python application.
Please see the [Python integration docs](https://posthog.com/docs/integrations/python-integration) for details.
## Installation
```bash
pip install hanzo-insights
```
## Quick Start
```python
from hanzo_insights import Insights
client = Insights('<your_project_api_key>', host='https://insights.hanzo.ai')
# Capture an event
client.capture('user_123', 'purchase', properties={'product': 'widget'})
# Feature flags
if client.feature_enabled('new-checkout', 'user_123'):
show_new_checkout()
```
## Module-level usage
```python
import hanzo_insights
hanzo_insights.api_key = '<your_project_api_key>'
hanzo_insights.host = 'https://insights.hanzo.ai'
hanzo_insights.capture('movie_played', distinct_id='user_123', properties={'movie_id': '42'})
hanzo_insights.shutdown()
```
## Python Version Support
| SDK Version | Python Versions Supported |
| -------------- | ----------------------------- |
| 7.3.1+ | 3.10, 3.11, 3.12, 3.13, 3.14 |
| 7.0.0 - 7.0.1 | 3.10, 3.11, 3.12, 3.13 |
| 4.0.1 - 6.x | 3.9, 3.10, 3.11, 3.12, 3.13 |
## Development
### Testing Locally
We recommend using [uv](https://docs.astral.sh/uv/). It's super fast.
1. Run `uv venv env` (creates virtual environment called "env")
* or `python3 -m venv env`
2. Run `source env/bin/activate` (activates the virtual environment)
3. Run `uv sync --extra dev --extra test` (installs the package in develop mode, along with test dependencies)
* or `pip install -e ".[dev,test]"`
4. you have to run `pre-commit install` to have auto linting pre commit
5. Run `make test`
1. To run a specific test do `pytest -k test_no_api_key`
## PostHog recommends `uv` so...
We use [uv](https://docs.astral.sh/uv/).
```bash
uv python install 3.9.19
uv python pin 3.9.19
uv python install 3.12
uv python pin 3.12
uv venv
source env/bin/activate
uv sync --extra dev --extra test
@@ -39,28 +57,14 @@ pre-commit install
make test
```
### Running Locally
### Running Tests
Assuming you have a [local version of PostHog](https://posthog.com/docs/developing-locally) running, you can run `python3 example.py` to see the library in action.
### Releasing Versions
Updates are released automatically using GitHub Actions when `version.py` is updated on `master`. After bumping `version.py` in `master` and adding to `CHANGELOG.md`, the [release workflow](https://github.com/PostHog/posthog-python/blob/master/.github/workflows/release.yaml) will automatically trigger and deploy the new version.
If you need to check the latest runs or manually trigger a release, you can go to [our release workflow's page](https://github.com/PostHog/posthog-python/actions/workflows/release.yaml) and dispatch it manually, using workflow from `master`.
### Testing changes locally with the PostHog app
You can run `make prep_local`, and it'll create a new folder alongside the SDK repo one called `posthog-python-local`, which you can then import into the posthog project by changing pyproject.toml to look like this:
```toml
dependencies = [
...
"posthoganalytics" #NOTE: no version number
...
]
...
[tools.uv.sources]
posthoganalytics = { path = "../posthog-python-local" }
```bash
make test
# or run a specific test:
pytest -k test_no_api_key
```
This'll let you build and test SDK changes fully locally, incorporating them into your local posthog app stack. It mainly takes care of the `posthog -> posthoganalytics` module renaming. You'll need to re-run `make prep_local` each time you make a change, and re-run `uv sync --active` in the posthog app project.
## License
MIT
+11
View File
@@ -0,0 +1,11 @@
# posthoganalytics
> **Do not use this package.** Use [`posthog`](https://pypi.org/project/posthog/) instead.
```bash
pip install posthog
```
This package exists solely for internal use by [posthog/posthog](https://github.com/posthog/posthog) to avoid import conflicts with the local `posthog` package in that repository. It is an automatically generated mirror of `posthog` — same code, same versions, just published under a different name.
If you are not working on the PostHog main repository, you should never need this package. All documentation, issues, and development happen in [`posthog-python`](https://github.com/posthog/posthog-python).
+1 -1
View File
@@ -4,5 +4,5 @@
source bin/helpers/_utils.sh
set_source_and_root_dir
flake8 posthog --ignore E501,W503
flake8 hanzo_insights --ignore E501,W503
mypy --no-site-packages --config-file mypy.ini . | mypy-baseline filter
+1 -1
View File
@@ -1,6 +1,6 @@
#!/usr/bin/env bash
#/ Usage: bin/docs
#/ Description: Generate documentation for the PostHog Python SDK
#/ Description: Generate documentation for the Insights Python SDK
source bin/helpers/_utils.sh
set_source_and_root_dir
ensure_virtual_env
+6 -6
View File
@@ -1,15 +1,15 @@
"""
Constants for PostHog Python SDK documentation generation.
Constants for Insights Python SDK documentation generation.
"""
from typing import Dict, Union
from posthog.version import VERSION
from hanzo_insights.version import VERSION
# Documentation generation metadata
DOCUMENTATION_METADATA = {
"hogRef": "0.3",
"slugPrefix": "posthog-python",
"specUrl": "https://github.com/PostHog/posthog-python",
"slugPrefix": "insights-python",
"specUrl": "https://github.com/Insights/insights-python",
}
# Docstring parsing patterns for new format
@@ -29,8 +29,8 @@ DOCSTRING_PATTERNS = {
# Output file configuration
OUTPUT_CONFIG: Dict[str, Union[str, int]] = {
"output_dir": "./references",
"filename": f"posthog-python-references-{VERSION}.json",
"filename_latest": "posthog-python-references-latest.json",
"filename": f"insights-python-references-{VERSION}.json",
"filename_latest": "insights-python-references-latest.json",
"indent": 2,
}
+26 -26
View File
@@ -1,6 +1,6 @@
#!/usr/bin/env python3
"""
Generate comprehensive SDK documentation JSON from PostHog Python SDK.
Generate comprehensive SDK documentation JSON from Insights Python SDK.
This script inspects the code and docstrings to create documentation in the specified format.
"""
@@ -337,19 +337,19 @@ def analyze_type(cls) -> dict:
def generate_sdk_documentation():
"""Generate complete SDK documentation in the requested format."""
# Import PostHog components
import posthog
from posthog.client import Client
import posthog.types as types_module
import posthog.args as args_module
from posthog.version import VERSION
# Import Insights components
import hanzo_insights
from hanzo_insights.client import Client
import hanzo_insights.types as types_module
import hanzo_insights.args as args_module
from hanzo_insights.version import VERSION
# Main SDK info
sdk_info = {
"version": VERSION,
"id": "posthog-python",
"title": "PostHog Python SDK",
"description": "Integrate PostHog into any python application.",
"id": "insights-python",
"title": "Insights Python SDK",
"description": "Integrate Insights into any python application.",
"slugPrefix": DOCUMENTATION_METADATA["slugPrefix"],
"specUrl": DOCUMENTATION_METADATA["specUrl"],
}
@@ -357,7 +357,7 @@ def generate_sdk_documentation():
# Collect types
types_list = []
# Types from posthog.types
# Types from hanzo_insights.types
for name in dir(types_module):
obj = getattr(types_module, name)
if inspect.isclass(obj) and not name.startswith("_"):
@@ -367,7 +367,7 @@ def generate_sdk_documentation():
except Exception as e:
print(f"Error analyzing type {name}: {e}")
# Types from posthog.args
# Types from hanzo_insights.args
for name in dir(args_module):
obj = getattr(args_module, name)
if inspect.isclass(obj) and not name.startswith("_"):
@@ -388,26 +388,26 @@ def generate_sdk_documentation():
# Collect classes
classes_list = []
# Main PostHog class (renamed from Client)
# Main Insights class (renamed from Client)
client_class = analyze_class(Client)
client_class["id"] = "PostHog"
client_class["title"] = "PostHog"
client_class["id"] = "Insights"
client_class["title"] = "Insights"
classes_list.append(client_class)
# Global module functions (functions callable as posthog.function_name)
# Global module functions (functions callable as hanzo_insights.function_name)
global_functions = []
for func_name in dir(posthog):
for func_name in dir(hanzo_insights):
# Skip private functions and non-callables
if func_name.startswith("_") or not callable(getattr(posthog, func_name)):
if func_name.startswith("_") or not callable(getattr(hanzo_insights, func_name)):
continue
func = getattr(posthog, func_name)
# Only include functions actually defined in the posthog module (not imported)
func = getattr(hanzo_insights, func_name)
# Only include functions actually defined in the hanzo_insights module (not imported)
# and exclude class references
if (
func_name not in ["Client", "Posthog"]
func_name not in ["Client", "Insights"]
and hasattr(func, "__module__")
and func.__module__ == "posthog"
and func.__module__ == "hanzo_insights"
):
try:
func_info = analyze_function(func, func_name)
@@ -421,8 +421,8 @@ def generate_sdk_documentation():
classes_list.append(
{
"id": "PostHogModule",
"title": "PostHog Module Functions",
"description": "Global functions available in the PostHog module",
"title": "Insights Module Functions",
"description": "Global functions available in the Insights module",
"functions": global_functions,
}
)
@@ -443,7 +443,7 @@ def generate_sdk_documentation():
# Create the final structure
result = {
"id": "posthog-python",
"id": "insights-python",
"hogRef": DOCUMENTATION_METADATA["hogRef"],
"info": sdk_info,
"types": types_list,
@@ -455,7 +455,7 @@ def generate_sdk_documentation():
if __name__ == "__main__":
print("Generating PostHog Python SDK documentation...")
print("Generating Insights Python SDK documentation...")
try:
documentation = generate_sdk_documentation()
+136 -130
View File
@@ -1,18 +1,18 @@
# PostHog Python library example
# Hanzo Insights Python library example
#
# This script demonstrates various PostHog Python SDK capabilities including:
# This script demonstrates various Hanzo Insights Python SDK capabilities including:
# - Basic event capture and user identification
# - Feature flag local evaluation
# - Feature flag payloads
# - Context management and tagging
#
# Setup:
# 1. Copy .env.example to .env and fill in your PostHog credentials
# 1. Copy .env.example to .env and fill in your Insights credentials
# 2. Run this script and choose from the interactive menu
import os
import posthog
import hanzo_insights
def load_env_file():
@@ -31,58 +31,44 @@ def load_env_file():
load_env_file()
# Get configuration
project_key = os.getenv("POSTHOG_PROJECT_API_KEY", "")
personal_api_key = os.getenv("POSTHOG_PERSONAL_API_KEY", "")
host = os.getenv("POSTHOG_HOST", "http://localhost:8000")
project_key = os.getenv("INSIGHTS_PROJECT_API_KEY", "")
personal_api_key = os.getenv("INSIGHTS_PERSONAL_API_KEY", "")
host = os.getenv("INSIGHTS_HOST", "http://localhost:8000")
# Check if credentials are provided
if not project_key or not personal_api_key:
print("❌ Missing PostHog credentials!")
print(
" Please set POSTHOG_PROJECT_API_KEY and POSTHOG_PERSONAL_API_KEY environment variables"
)
# Check if project key is provided (required)
if not project_key:
print("❌ Missing Insights project API key!")
print(" Please set INSIGHTS_PROJECT_API_KEY environment variable")
print(" or copy .env.example to .env and fill in your values")
exit(1)
# Test authentication before proceeding
print("🔑 Testing PostHog authentication...")
# Configure Insights with credentials
hanzo_insights.debug = False
hanzo_insights.api_key = project_key
hanzo_insights.project_api_key = project_key
hanzo_insights.host = host
hanzo_insights.poll_interval = 10
try:
# Configure PostHog with credentials
posthog.debug = False # Keep quiet during auth test
posthog.api_key = project_key
posthog.project_api_key = project_key
posthog.personal_api_key = personal_api_key
posthog.host = host
posthog.poll_interval = 10
# Check if personal API key is available for local evaluation
local_eval_available = bool(personal_api_key)
if personal_api_key:
hanzo_insights.personal_api_key = personal_api_key
# Test by attempting to get feature flags (this validates both keys)
# This will fail if credentials are invalid
test_flags = posthog.get_all_flags("test_user", only_evaluate_locally=True)
# If we get here without exception, credentials work
print("✅ Authentication successful!")
print(f" Project API Key: {project_key[:9]}...")
print(" Personal API Key: [REDACTED]")
print(f" Host: {host}\n\n")
except Exception as e:
print("❌ Authentication failed!")
print(f" Error: {e}")
print("\n Please check your credentials:")
print(" - POSTHOG_PROJECT_API_KEY: Project API key from PostHog settings")
print(
" - POSTHOG_PERSONAL_API_KEY: Personal API key (required for local evaluation)"
)
print(" - POSTHOG_HOST: Your PostHog instance URL")
exit(1)
print("🔑 Insights Configuration:")
print(f" Project API Key: {project_key[:9]}...")
if local_eval_available:
print(" Personal API Key: [SET]")
else:
print(" Personal API Key: [NOT SET] - Local evaluation examples will be skipped")
print(f" Host: {host}\n")
# Display menu and get user choice
print("🚀 PostHog Python SDK Demo - Choose an example to run:\n")
print("🚀 Hanzo Insights Python SDK Demo - Choose an example to run:\n")
print("1. Identify and capture examples")
print("2. Feature flag local evaluation examples")
local_eval_note = "" if local_eval_available else " [requires personal API key]"
print(f"2. Feature flag local evaluation examples{local_eval_note}")
print("3. Feature flag payload examples")
print("4. Flag dependencies examples")
print(f"4. Flag dependencies examples{local_eval_note}")
print("5. Context management and tagging examples")
print("6. Run all examples")
print("7. Exit")
@@ -93,11 +79,11 @@ if choice == "1":
print("IDENTIFY AND CAPTURE EXAMPLES")
print("=" * 60)
posthog.debug = True
hanzo_insights.debug = True
# Capture an event
print("📊 Capturing events...")
posthog.capture(
hanzo_insights.capture(
"event",
distinct_id="distinct_id",
properties={"property1": "value", "property2": "value"},
@@ -106,14 +92,14 @@ if choice == "1":
# Alias a previous distinct id with a new one
print("🔗 Creating alias...")
posthog.alias("distinct_id", "new_distinct_id")
hanzo_insights.alias("distinct_id", "new_distinct_id")
posthog.capture(
hanzo_insights.capture(
"event2",
distinct_id="new_distinct_id",
properties={"property1": "value", "property2": "value"},
)
posthog.capture(
hanzo_insights.capture(
"event-with-groups",
distinct_id="new_distinct_id",
properties={"property1": "value", "property2": "value"},
@@ -122,66 +108,74 @@ if choice == "1":
# Add properties to the person
print("👤 Identifying user...")
posthog.set(
hanzo_insights.set(
distinct_id="new_distinct_id", properties={"email": "something@something.com"}
)
# Add properties to a group
print("🏢 Identifying group...")
posthog.group_identify("company", "id:5", {"employees": 11})
hanzo_insights.group_identify("company", "id:5", {"employees": 11})
# Properties set only once to the person
print("🔒 Setting properties once...")
posthog.set_once(
hanzo_insights.set_once(
distinct_id="new_distinct_id", properties={"self_serve_signup": True}
)
# This will not change the property (because it was already set)
posthog.set_once(
hanzo_insights.set_once(
distinct_id="new_distinct_id", properties={"self_serve_signup": False}
)
print("🔄 Updating properties...")
posthog.set(distinct_id="new_distinct_id", properties={"current_browser": "Chrome"})
posthog.set(
hanzo_insights.set(distinct_id="new_distinct_id", properties={"current_browser": "Chrome"})
hanzo_insights.set(
distinct_id="new_distinct_id", properties={"current_browser": "Firefox"}
)
elif choice == "2":
if not local_eval_available:
print("\n❌ This example requires a personal API key for local evaluation.")
print(
" Set INSIGHTS_PERSONAL_API_KEY environment variable to run this example."
)
hanzo_insights.shutdown()
exit(1)
print("\n" + "=" * 60)
print("FEATURE FLAG LOCAL EVALUATION EXAMPLES")
print("=" * 60)
posthog.debug = True
hanzo_insights.debug = True
print("🏁 Testing basic feature flags...")
print(
f"beta-feature for 'distinct_id': {posthog.feature_enabled('beta-feature', 'distinct_id')}"
f"beta-feature for 'distinct_id': {hanzo_insights.feature_enabled('beta-feature', 'distinct_id')}"
)
print(
f"beta-feature for 'new_distinct_id': {posthog.feature_enabled('beta-feature', 'new_distinct_id')}"
f"beta-feature for 'new_distinct_id': {hanzo_insights.feature_enabled('beta-feature', 'new_distinct_id')}"
)
print(
f"beta-feature with groups: {posthog.feature_enabled('beta-feature-groups', 'distinct_id', groups={'company': 'id:5'})}"
f"beta-feature with groups: {hanzo_insights.feature_enabled('beta-feature-groups', 'distinct_id', groups={'company': 'id:5'})}"
)
print("\n🌍 Testing location-based flags...")
# Assume test-flag has `City Name = Sydney` as a person property set
print(
f"Sydney user: {posthog.feature_enabled('test-flag', 'random_id_12345', person_properties={'$geoip_city_name': 'Sydney'})}"
f"Sydney user: {hanzo_insights.feature_enabled('test-flag', 'random_id_12345', person_properties={'$geoip_city_name': 'Sydney'})}"
)
print(
f"Sydney user (local only): {posthog.feature_enabled('test-flag', 'distinct_id_random_22', person_properties={'$geoip_city_name': 'Sydney'}, only_evaluate_locally=True)}"
f"Sydney user (local only): {hanzo_insights.feature_enabled('test-flag', 'distinct_id_random_22', person_properties={'$geoip_city_name': 'Sydney'}, only_evaluate_locally=True)}"
)
print("\n📋 Getting all flags...")
print(f"All flags: {posthog.get_all_flags('distinct_id_random_22')}")
print(f"All flags: {hanzo_insights.get_all_flags('distinct_id_random_22')}")
print(
f"All flags (local): {posthog.get_all_flags('distinct_id_random_22', only_evaluate_locally=True)}"
f"All flags (local): {hanzo_insights.get_all_flags('distinct_id_random_22', only_evaluate_locally=True)}"
)
print(
f"All flags with properties: {posthog.get_all_flags('distinct_id_random_22', person_properties={'$geoip_city_name': 'Sydney'}, only_evaluate_locally=True)}"
f"All flags with properties: {hanzo_insights.get_all_flags('distinct_id_random_22', person_properties={'$geoip_city_name': 'Sydney'}, only_evaluate_locally=True)}"
)
elif choice == "3":
@@ -189,22 +183,22 @@ elif choice == "3":
print("FEATURE FLAG PAYLOAD EXAMPLES")
print("=" * 60)
posthog.debug = True
hanzo_insights.debug = True
print("📦 Testing feature flag payloads...")
print(
f"beta-feature payload: {posthog.get_feature_flag_payload('beta-feature', 'distinct_id')}"
f"beta-feature payload: {hanzo_insights.get_feature_flag_payload('beta-feature', 'distinct_id')}"
)
print(
f"All flags and payloads: {posthog.get_all_flags_and_payloads('distinct_id')}"
f"All flags and payloads: {hanzo_insights.get_all_flags_and_payloads('distinct_id')}"
)
print(
f"Remote config payload: {posthog.get_remote_config_payload('encrypted_payload_flag_key')}"
f"Remote config payload: {hanzo_insights.get_remote_config_payload('encrypted_payload_flag_key')}"
)
# Get feature flag result with all details (enabled, variant, payload, key, reason)
print("\n🔍 Getting detailed flag result...")
result = posthog.get_feature_flag_result("beta-feature", "distinct_id")
result = hanzo_insights.get_feature_flag_result("beta-feature", "distinct_id")
if result:
print(f"Flag key: {result.key}")
print(f"Flag enabled: {result.enabled}")
@@ -215,6 +209,14 @@ elif choice == "3":
print(f"Value (variant or enabled): {result.get_value()}")
elif choice == "4":
if not local_eval_available:
print("\n❌ This example requires a personal API key for local evaluation.")
print(
" Set INSIGHTS_PERSONAL_API_KEY environment variable to run this example."
)
hanzo_insights.shutdown()
exit(1)
print("\n" + "=" * 60)
print("FLAG DEPENDENCIES EXAMPLES")
print("=" * 60)
@@ -232,10 +234,10 @@ elif choice == "4":
print(" - Rollout: 100%")
print("")
posthog.debug = True
hanzo_insights.debug = True
# Test @example.com user (should satisfy dependency if flags exist)
result1 = posthog.feature_enabled(
result1 = hanzo_insights.feature_enabled(
"test-flag-dependency",
"example_user",
person_properties={"email": "user@example.com"},
@@ -244,7 +246,7 @@ elif choice == "4":
print(f"✅ @example.com user (test-flag-dependency): {result1}")
# Test non-example.com user (dependency should not be satisfied)
result2 = posthog.feature_enabled(
result2 = hanzo_insights.feature_enabled(
"test-flag-dependency",
"regular_user",
person_properties={"email": "user@other.com"},
@@ -253,13 +255,13 @@ elif choice == "4":
print(f"❌ Regular user (test-flag-dependency): {result2}")
# Test beta-feature directly for comparison
beta1 = posthog.feature_enabled(
beta1 = hanzo_insights.feature_enabled(
"beta-feature",
"example_user",
person_properties={"email": "user@example.com"},
only_evaluate_locally=True,
)
beta2 = posthog.feature_enabled(
beta2 = hanzo_insights.feature_enabled(
"beta-feature",
"regular_user",
person_properties={"email": "user@other.com"},
@@ -301,7 +303,7 @@ elif choice == "4":
print("")
# Test pineapple -> blue -> breaking-bad chain
dependent_result3 = posthog.get_feature_flag(
dependent_result3 = hanzo_insights.get_feature_flag(
"multivariate-root-flag",
"regular_user",
person_properties={"email": "pineapple@example.com"},
@@ -315,7 +317,7 @@ elif choice == "4":
print("'multivariate-root-flag' with email pineapple@example.com succeeded")
# Test mango -> red -> the-wire chain
dependent_result4 = posthog.get_feature_flag(
dependent_result4 = hanzo_insights.get_feature_flag(
"multivariate-root-flag",
"regular_user",
person_properties={"email": "mango@example.com"},
@@ -334,19 +336,19 @@ elif choice == "4":
("pineapple@example.com", ["pineapple", "blue", "breaking-bad"]),
("mango@example.com", ["mango", "red", "the-wire"]),
]:
leaf = posthog.get_feature_flag(
leaf = hanzo_insights.get_feature_flag(
"multivariate-leaf-flag",
"regular_user",
person_properties={"email": email},
only_evaluate_locally=True,
)
intermediate = posthog.get_feature_flag(
intermediate = hanzo_insights.get_feature_flag(
"multivariate-intermediate-flag",
"regular_user",
person_properties={"email": email},
only_evaluate_locally=True,
)
root = posthog.get_feature_flag(
root = hanzo_insights.get_feature_flag(
"multivariate-root-flag",
"regular_user",
person_properties={"email": email},
@@ -371,7 +373,7 @@ elif choice == "5":
print("CONTEXT MANAGEMENT AND TAGGING EXAMPLES")
print("=" * 60)
posthog.debug = True
hanzo_insights.debug = True
print("🏷️ Testing context management...")
print(
@@ -382,12 +384,12 @@ elif choice == "5":
# and tagged with the context tags. Other events captured will also be tagged with the context tags. By default,
# the new context inherits tags from the parent context.
try:
with posthog.new_context():
posthog.tag("transaction_id", "abc123")
posthog.tag("some_arbitrary_value", {"tags": "can be dicts"})
with hanzo_insights.new_context():
hanzo_insights.tag("transaction_id", "abc123")
hanzo_insights.tag("some_arbitrary_value", {"tags": "can be dicts"})
# This event will be captured with the tags set above
posthog.capture("order_processed")
hanzo_insights.capture("order_processed")
print("✅ Event captured with inherited context tags")
# This exception will be captured with the tags set above
# raise Exception("Order processing failed")
@@ -396,30 +398,30 @@ elif choice == "5":
# Use fresh=True to start with a clean context (no inherited tags)
try:
with posthog.new_context(fresh=True):
posthog.tag("session_id", "xyz789")
with hanzo_insights.new_context(fresh=True):
hanzo_insights.tag("session_id", "xyz789")
# Only session_id tag will be present, no inherited tags
posthog.capture("session_event")
hanzo_insights.capture("session_event")
print("✅ Event captured with fresh context tags")
# raise Exception("Session handling failed")
except Exception as e:
print(f"Exception captured: {e}")
# You can also use the `@posthog.scoped()` decorator to enter a new context.
# You can also use the `@hanzo_insights.scoped()` decorator to enter a new context.
# By default, it inherits tags from the parent context
@posthog.scoped()
@hanzo_insights.scoped()
def process_order(order_id):
posthog.tag("order_id", order_id)
posthog.capture("order_step_completed")
hanzo_insights.tag("order_id", order_id)
hanzo_insights.capture("order_step_completed")
print(f"✅ Order {order_id} processed with scoped context")
# Exception will be captured and tagged automatically
# raise Exception("Order processing failed")
# Use fresh=True to start with a clean context (no inherited tags)
@posthog.scoped(fresh=True)
@hanzo_insights.scoped(fresh=True)
def process_payment(payment_id):
posthog.tag("payment_id", payment_id)
posthog.capture("payment_processed")
hanzo_insights.tag("payment_id", payment_id)
hanzo_insights.capture("payment_processed")
print(f"✅ Payment {payment_id} processed with fresh scoped context")
# Only payment_id tag will be present, no inherited tags
# raise Exception("Payment processing failed")
@@ -429,74 +431,78 @@ elif choice == "5":
elif choice == "6":
print("\n🔄 Running all examples...")
if not local_eval_available:
print(" (Skipping local evaluation examples - no personal API key set)\n")
# Run example 1
print(f"\n{'🔸' * 20} IDENTIFY AND CAPTURE {'🔸' * 20}")
posthog.debug = True
hanzo_insights.debug = True
print("📊 Capturing events...")
posthog.capture(
hanzo_insights.capture(
"event",
distinct_id="distinct_id",
properties={"property1": "value", "property2": "value"},
send_feature_flags=True,
)
print("🔗 Creating alias...")
posthog.alias("distinct_id", "new_distinct_id")
hanzo_insights.alias("distinct_id", "new_distinct_id")
print("👤 Identifying user...")
posthog.set(
hanzo_insights.set(
distinct_id="new_distinct_id", properties={"email": "something@something.com"}
)
# Run example 2
print(f"\n{'🔸' * 20} FEATURE FLAGS {'🔸' * 20}")
print("🏁 Testing basic feature flags...")
print(f"beta-feature: {posthog.feature_enabled('beta-feature', 'distinct_id')}")
print(
f"Sydney user: {posthog.feature_enabled('test-flag', 'random_id_12345', person_properties={'$geoip_city_name': 'Sydney'})}"
)
# Run example 2 (requires local evaluation)
if local_eval_available:
print(f"\n{'🔸' * 20} FEATURE FLAGS {'🔸' * 20}")
print("🏁 Testing basic feature flags...")
print(f"beta-feature: {hanzo_insights.feature_enabled('beta-feature', 'distinct_id')}")
print(
f"Sydney user: {hanzo_insights.feature_enabled('test-flag', 'random_id_12345', person_properties={'$geoip_city_name': 'Sydney'})}"
)
# Run example 3
print(f"\n{'🔸' * 20} PAYLOADS {'🔸' * 20}")
print("📦 Testing payloads...")
print(f"Payload: {posthog.get_feature_flag_payload('beta-feature', 'distinct_id')}")
print(f"Payload: {hanzo_insights.get_feature_flag_payload('beta-feature', 'distinct_id')}")
# Run example 4
print(f"\n{'🔸' * 20} FLAG DEPENDENCIES {'🔸' * 20}")
print("🔗 Testing flag dependencies...")
result1 = posthog.feature_enabled(
"test-flag-dependency",
"demo_user",
person_properties={"email": "user@example.com"},
only_evaluate_locally=True,
)
result2 = posthog.feature_enabled(
"test-flag-dependency",
"demo_user2",
person_properties={"email": "user@other.com"},
only_evaluate_locally=True,
)
print(f"✅ @example.com user: {result1}, regular user: {result2}")
# Run example 4 (requires local evaluation)
if local_eval_available:
print(f"\n{'🔸' * 20} FLAG DEPENDENCIES {'🔸' * 20}")
print("🔗 Testing flag dependencies...")
result1 = hanzo_insights.feature_enabled(
"test-flag-dependency",
"demo_user",
person_properties={"email": "user@example.com"},
only_evaluate_locally=True,
)
result2 = hanzo_insights.feature_enabled(
"test-flag-dependency",
"demo_user2",
person_properties={"email": "user@other.com"},
only_evaluate_locally=True,
)
print(f"✅ @example.com user: {result1}, regular user: {result2}")
# Run example 5
print(f"\n{'🔸' * 20} CONTEXT MANAGEMENT {'🔸' * 20}")
print("🏷️ Testing context management...")
with posthog.new_context():
posthog.tag("demo_run", "all_examples")
posthog.capture("demo_completed")
with hanzo_insights.new_context():
hanzo_insights.tag("demo_run", "all_examples")
hanzo_insights.capture("demo_completed")
print("✅ Demo completed with context tags")
elif choice == "7":
print("👋 Goodbye!")
posthog.shutdown()
hanzo_insights.shutdown()
exit()
else:
print("❌ Invalid choice. Please run again and select 1-7.")
posthog.shutdown()
hanzo_insights.shutdown()
exit()
print("\n" + "=" * 60)
print("✅ Example completed!")
print("=" * 60)
posthog.shutdown()
hanzo_insights.shutdown()
+144
View File
@@ -0,0 +1,144 @@
"""
Redis-based distributed cache for Insights feature flag definitions.
This example demonstrates how to implement a FlagDefinitionCacheProvider
using Redis for multi-instance deployments (leader election pattern).
Usage:
import redis
from hanzo_insights import Insights
redis_client = redis.Redis(host='localhost', port=6379, decode_responses=True)
cache = RedisFlagCache(redis_client, service_key="my-service")
client = Insights(
"<project_api_key>",
personal_api_key="<personal_api_key>",
flag_definition_cache_provider=cache,
)
Requirements:
pip install redis
"""
import json
import uuid
from hanzo_insights import FlagDefinitionCacheData, FlagDefinitionCacheProvider
from redis import Redis
from typing import Optional
class RedisFlagCache(FlagDefinitionCacheProvider):
"""
A distributed cache for Insights feature flag definitions using Redis.
In a multi-instance deployment (e.g., multiple serverless functions or containers),
we want only ONE instance to poll Insights for flag updates, while all instances
share the cached results. This prevents N instances from making N redundant API calls.
The implementation uses leader election:
- One instance "wins" and becomes responsible for fetching
- Other instances read from the shared cache
- If the leader dies, the lock expires (TTL) and another instance takes over
Uses Lua scripts for atomic operations, following Redis distributed lock best practices:
https://redis.io/docs/latest/develop/clients/patterns/distributed-locks/
"""
LOCK_TTL_MS = 60 * 1000 # 60 seconds, should be longer than the flags poll interval
CACHE_TTL_SECONDS = 60 * 60 * 24 # 24 hours
# Lua script: acquire lock if free, or extend if we own it
_LUA_TRY_LEAD = """
local current = redis.call('GET', KEYS[1])
if current == false then
redis.call('SET', KEYS[1], ARGV[1], 'PX', ARGV[2])
return 1
elseif current == ARGV[1] then
redis.call('PEXPIRE', KEYS[1], ARGV[2])
return 1
end
return 0
"""
# Lua script: release lock only if we own it
_LUA_STOP_LEAD = """
if redis.call('GET', KEYS[1]) == ARGV[1] then
return redis.call('DEL', KEYS[1])
end
return 0
"""
def __init__(self, redis: Redis[str], service_key: str):
"""
Initialize the Redis flag cache.
Args:
redis: A redis-py client instance. Must be configured with
decode_responses=True for correct string handling.
service_key: A unique identifier for this service/environment.
Used to scope Redis keys, allowing multiple services
or environments to share the same Redis instance.
Examples: "my-api-prod", "checkout-service", "staging".
Redis Keys Created:
- insights:flags:{service_key} - Cached flag definitions (JSON)
- insights:flags:{service_key}:lock - Leader election lock
Example:
redis_client = redis.Redis(
host='localhost',
port=6379,
decode_responses=True
)
cache = RedisFlagCache(redis_client, service_key="my-api-prod")
"""
self._redis = redis
self._cache_key = f"insights:flags:{service_key}"
self._lock_key = f"insights:flags:{service_key}:lock"
self._instance_id = str(uuid.uuid4())
self._try_lead = self._redis.register_script(self._LUA_TRY_LEAD)
self._stop_lead = self._redis.register_script(self._LUA_STOP_LEAD)
def get_flag_definitions(self) -> Optional[FlagDefinitionCacheData]:
"""
Retrieve cached flag definitions from Redis.
Returns:
Cached flag definitions if available, None otherwise.
"""
cached = self._redis.get(self._cache_key)
return json.loads(cached) if cached else None
def should_fetch_flag_definitions(self) -> bool:
"""
Determines if this instance should fetch flag definitions from Insights.
Atomically either:
- Acquires the lock if no one holds it, OR
- Extends the lock TTL if we already hold it
Returns:
True if this instance is the leader and should fetch, False otherwise.
"""
result = self._try_lead(
keys=[self._lock_key],
args=[self._instance_id, self.LOCK_TTL_MS],
)
return result == 1
def on_flag_definitions_received(self, data: FlagDefinitionCacheData) -> None:
"""
Store fetched flag definitions in Redis.
Args:
data: The flag definitions to cache.
"""
self._redis.set(self._cache_key, json.dumps(data), ex=self.CACHE_TTL_SECONDS)
def shutdown(self) -> None:
"""
Release leadership if we hold it. Safe to call even if not the leader.
"""
self._stop_lead(keys=[self._lock_key], args=[self._instance_id])
@@ -1,15 +1,15 @@
#!/usr/bin/env python3
"""
Simple test script for PostHog remote config endpoint.
Simple test script for Insights remote config endpoint.
"""
import posthog
import hanzo_insights
# Initialize PostHog client
posthog.api_key = "phc_..."
posthog.personal_api_key = "phs_..." # or "phx_..."
posthog.host = "http://localhost:8000" # or "https://us.posthog.com"
posthog.debug = True
# Initialize Insights client
hanzo_insights.api_key = "phc_..."
hanzo_insights.personal_api_key = "phs_..." # or "phx_..."
hanzo_insights.host = "http://localhost:8000" # or "https://us.insights.hanzo.ai"
hanzo_insights.debug = True
def test_remote_config():
@@ -21,7 +21,7 @@ def test_remote_config():
try:
# Get remote config payload
payload = posthog.get_remote_config_payload(flag_key)
payload = hanzo_insights.get_remote_config_payload(flag_key)
print(f"✅ Success! Remote config payload for '{flag_key}': {payload}")
except Exception as e:
@@ -1,26 +1,68 @@
import datetime # noqa: F401
from typing import Callable, Dict, Optional, Any # noqa: F401
from typing import Any, Callable, Dict, Optional # noqa: F401
from typing_extensions import Unpack
from posthog.args import OptionalCaptureArgs, OptionalSetArgs, ExceptionArg
from posthog.client import Client
from posthog.contexts import (
new_context as inner_new_context,
scoped as inner_scoped,
tag as inner_tag,
set_context_session as inner_set_context_session,
from hanzo_insights.args import ExceptionArg, OptionalCaptureArgs, OptionalSetArgs
from hanzo_insights.client import Client
from hanzo_insights.contexts import (
identify_context as inner_identify_context,
)
from hanzo_insights.contexts import (
new_context as inner_new_context,
)
from hanzo_insights.contexts import (
scoped as inner_scoped,
)
from hanzo_insights.contexts import (
set_capture_exception_code_variables_context as inner_set_capture_exception_code_variables_context,
set_code_variables_mask_patterns_context as inner_set_code_variables_mask_patterns_context,
)
from hanzo_insights.contexts import (
set_code_variables_ignore_patterns_context as inner_set_code_variables_ignore_patterns_context,
)
from posthog.exception_utils import (
from hanzo_insights.contexts import (
set_code_variables_mask_patterns_context as inner_set_code_variables_mask_patterns_context,
)
from hanzo_insights.contexts import (
set_context_device_id as inner_set_context_device_id,
)
from hanzo_insights.contexts import (
set_context_session as inner_set_context_session,
)
from hanzo_insights.contexts import (
tag as inner_tag,
)
from hanzo_insights.contexts import (
get_tags as inner_get_tags,
)
from hanzo_insights.exception_utils import (
DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS,
DEFAULT_CODE_VARIABLES_MASK_PATTERNS,
)
from posthog.feature_flags import InconclusiveMatchError, RequiresServerEvaluation
from posthog.types import FeatureFlag, FlagsAndPayloads, FeatureFlagResult
from posthog.version import VERSION
from hanzo_insights.feature_flags import (
InconclusiveMatchError as InconclusiveMatchError,
)
from hanzo_insights.feature_flags import (
RequiresServerEvaluation as RequiresServerEvaluation,
)
from hanzo_insights.flag_definition_cache import (
FlagDefinitionCacheData as FlagDefinitionCacheData,
FlagDefinitionCacheProvider as FlagDefinitionCacheProvider,
)
from hanzo_insights.request import (
disable_connection_reuse as disable_connection_reuse,
enable_keep_alive as enable_keep_alive,
set_socket_options as set_socket_options,
SocketOptions as SocketOptions,
)
from hanzo_insights.types import (
FeatureFlag,
FlagsAndPayloads,
)
from hanzo_insights.types import (
FeatureFlagResult as FeatureFlagResult,
)
from hanzo_insights.version import VERSION
__version__ = VERSION
@@ -34,11 +76,11 @@ def new_context(fresh=False, capture_exceptions=True, client=None):
Args:
fresh: Whether to start with a fresh context (default: False)
capture_exceptions: Whether to capture exceptions raised within the context (default: True)
client: Optional Posthog client instance to use for this context (default: None)
client: Optional Insights client instance to use for this context (default: None)
Examples:
```python
from posthog import new_context, tag, capture
from hanzo_insights import new_context, tag, capture
with new_context():
tag("request_id", "123")
capture("event_name", properties={"property": "value"})
@@ -58,11 +100,11 @@ def scoped(fresh=False, capture_exceptions=True):
Args:
fresh: Whether to start with a fresh context (default: False)
capture_exceptions: Whether to capture and track exceptions with posthog error tracking (default: True)
capture_exceptions: Whether to capture and track exceptions with Insights error tracking (default: True)
Examples:
```python
from posthog import scoped, tag, capture
from hanzo_insights import scoped, tag, capture
@scoped()
def process_payment(payment_id):
tag("payment_id", payment_id)
@@ -84,7 +126,7 @@ def set_context_session(session_id: str):
Examples:
```python
from posthog import set_context_session
from hanzo_insights import set_context_session
set_context_session("session_123")
```
@@ -94,6 +136,26 @@ def set_context_session(session_id: str):
return inner_set_context_session(session_id)
def set_context_device_id(device_id: str):
"""
Set the device ID for the current context, associating all feature flag requests
in this or child contexts with the given device ID.
Args:
device_id: The device ID to associate with the current context and its children
Examples:
```python
from hanzo_insights import set_context_device_id
set_context_device_id("device_123")
```
Category:
Contexts
"""
return inner_set_context_device_id(device_id)
def identify_context(distinct_id: str):
"""
Identify the current context with a distinct ID.
@@ -103,7 +165,7 @@ def identify_context(distinct_id: str):
Examples:
```python
from posthog import identify_context
from hanzo_insights import identify_context
identify_context("user_123")
```
@@ -144,7 +206,7 @@ def tag(name: str, value: Any):
Examples:
```python
from posthog import tag
from hanzo_insights import tag
tag("user_id", "123")
```
@@ -154,6 +216,19 @@ def tag(name: str, value: Any):
return inner_tag(name, value)
def get_tags() -> Dict[str, Any]:
"""
Get all tags from the current context.
Returns:
Dict of all tags in the current context
Category:
Contexts
"""
return inner_get_tags()
"""Settings."""
api_key = None # type: Optional[str]
host = None # type: Optional[str]
@@ -184,10 +259,11 @@ default_client = None # type: Optional[Client]
capture_exception_code_variables = False
code_variables_mask_patterns = DEFAULT_CODE_VARIABLES_MASK_PATTERNS
code_variables_ignore_patterns = DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS
in_app_modules = None # type: Optional[list[str]]
# NOTE - this and following functions take unpacked kwargs because we needed to make
# it impossible to write `posthog.capture(distinct-id, event-name)` - basically, to enforce
# it impossible to write `hanzo_insights.capture(distinct-id, event-name)` - basically, to enforce
# the breaking change made between 5.3.0 and 6.0.0. This decision can be unrolled in later
# versions, without a breaking change, to get back the type information in function signatures
def capture(event: str, **kwargs: Unpack[OptionalCaptureArgs]) -> Optional[str]:
@@ -204,12 +280,12 @@ def capture(event: str, **kwargs: Unpack[OptionalCaptureArgs]) -> Optional[str]:
disable_geoip: Whether to disable GeoIP lookup
Details:
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. A capture call requires an event name to specify the event. We recommend using [verb] [noun], like `movie played` or `movie updated` to easily identify what your events mean later on. Capture takes a number of optional arguments, which are defined by the `OptionalCaptureArgs` type.
Capture allows you to capture anything a user does within your system, which you can later use in Insights to find patterns in usage, work out which features to improve or where people are giving up. A capture call requires an event name to specify the event. We recommend using [verb] [noun], like `movie played` or `movie updated` to easily identify what your events mean later on. Capture takes a number of optional arguments, which are defined by the `OptionalCaptureArgs` type.
Examples:
```python
# Context and capture usage
from posthog import new_context, identify_context, tag_context, capture
from hanzo_insights import new_context, identify_context, tag_context, capture
# Enter a new context (e.g. a request/response cycle, an instance of a background job, etc)
with new_context():
# Associate this context with some user, by distinct_id
@@ -236,7 +312,7 @@ def capture(event: str, **kwargs: Unpack[OptionalCaptureArgs]) -> Optional[str]:
```
```python
# Set event properties
from posthog import capture
from hanzo_insights import capture
capture(
"user_signed_up",
distinct_id="distinct_id_of_the_user",
@@ -263,7 +339,7 @@ def set(**kwargs: Unpack[OptionalSetArgs]) -> Optional[str]:
Examples:
```python
# Set person properties
from posthog import capture
from hanzo_insights import capture
capture(
'distinct_id',
event='event_name',
@@ -290,7 +366,7 @@ def set_once(**kwargs: Unpack[OptionalSetArgs]) -> Optional[str]:
Examples:
```python
# Set property once
from posthog import capture
from hanzo_insights import capture
capture(
'distinct_id',
event='event_name',
@@ -330,7 +406,7 @@ def group_identify(
Examples:
```python
# Group identify
from posthog import group_identify
from hanzo_insights import group_identify
group_identify('company', 'company_id_in_your_db', {
'name': 'Awesome Inc.',
'employees': 11
@@ -375,7 +451,7 @@ def alias(
Examples:
```python
# Alias user
from posthog import alias
from hanzo_insights import alias
alias(previous_id='distinct_id', distinct_id='alias_id')
```
Category:
@@ -403,12 +479,12 @@ def capture_exception(
exception: The exception to capture. If not provided, the current exception is captured via `sys.exc_info()`
Details:
Capture exception is idempotent - if it is called twice with the same exception instance, only a occurrence will be tracked in posthog. This is because, generally, contexts will cause exceptions to be captured automatically. However, to ensure you track an exception, if you catch and do not re-raise it, capturing it manually is recommended, unless you are certain it will have crossed a context boundary (e.g. by existing a `with posthog.new_context():` block already). If the passed exception was raised and caught, the captured stack trace will consist of every frame between where the exception was raised and the point at which it is captured (the "traceback"). If the passed exception was never raised, e.g. if you call `posthog.capture_exception(ValueError("Some Error"))`, the stack trace captured will be the full stack trace at the moment the exception was captured. Note that heavy use of contexts will lead to truncated stack traces, as the exception will be captured by the context entered most recently, which may not be the point you catch the exception for the final time in your code. It's recommended to use contexts sparingly, for this reason. `capture_exception` takes the same set of optional arguments as `capture`.
Capture exception is idempotent - if it is called twice with the same exception instance, only a occurrence will be tracked in hanzo_insights. This is because, generally, contexts will cause exceptions to be captured automatically. However, to ensure you track an exception, if you catch and do not re-raise it, capturing it manually is recommended, unless you are certain it will have crossed a context boundary (e.g. by existing a `with hanzo_insights.new_context():` block already). If the passed exception was raised and caught, the captured stack trace will consist of every frame between where the exception was raised and the point at which it is captured (the "traceback"). If the passed exception was never raised, e.g. if you call `hanzo_insights.capture_exception(ValueError("Some Error"))`, the stack trace captured will be the full stack trace at the moment the exception was captured. Note that heavy use of contexts will lead to truncated stack traces, as the exception will be captured by the context entered most recently, which may not be the point you catch the exception for the final time in your code. It's recommended to use contexts sparingly, for this reason. `capture_exception` takes the same set of optional arguments as `capture`.
Examples:
```python
# Capture exception
from posthog import capture_exception
from hanzo_insights import capture_exception
try:
risky_operation()
except Exception as e:
@@ -430,6 +506,7 @@ def feature_enabled(
only_evaluate_locally=False, # type: bool
send_feature_flag_events=True, # type: bool
disable_geoip=None, # type: Optional[bool]
device_id=None, # type: Optional[str]
):
# type: (...) -> bool
"""
@@ -446,12 +523,12 @@ def feature_enabled(
disable_geoip: Whether to disable GeoIP lookup
Details:
You can call `posthog.load_feature_flags()` before to make sure you're not doing unexpected requests.
You can call `hanzo_insights.load_feature_flags()` before to make sure you're not doing unexpected requests.
Examples:
```python
# Boolean feature flag
from posthog import feature_enabled, get_feature_flag_payload
from hanzo_insights import feature_enabled, get_feature_flag_payload
is_my_flag_enabled = feature_enabled('flag-key', 'distinct_id_of_your_user')
if is_my_flag_enabled:
matched_flag_payload = get_feature_flag_payload('flag-key', 'distinct_id_of_your_user')
@@ -469,6 +546,7 @@ def feature_enabled(
only_evaluate_locally=only_evaluate_locally,
send_feature_flag_events=send_feature_flag_events,
disable_geoip=disable_geoip,
device_id=device_id,
)
@@ -481,6 +559,7 @@ def get_feature_flag(
only_evaluate_locally=False, # type: bool
send_feature_flag_events=True, # type: bool
disable_geoip=None, # type: Optional[bool]
device_id=None, # type: Optional[str]
) -> Optional[FeatureFlag]:
"""
Get feature flag variant for users. Used with experiments.
@@ -496,12 +575,12 @@ def get_feature_flag(
disable_geoip: Whether to disable GeoIP lookup
Details:
`groups` are a mapping from group type to group key. So, if you have a group type of "organization" and a group key of "5", you would pass groups={"organization": "5"}. `group_properties` take the format: { group_type_name: { group_properties } }. So, for example, if you have the group type "organization" and the group key "5", with the properties name, and employee count, you'll send these as: group_properties={"organization": {"name": "PostHog", "employees": 11}}.
`groups` are a mapping from group type to group key. So, if you have a group type of "organization" and a group key of "5", you would pass groups={"organization": "5"}. `group_properties` take the format: { group_type_name: { group_properties } }. So, for example, if you have the group type "organization" and the group key "5", with the properties name, and employee count, you'll send these as: group_properties={"organization": {"name": "Hanzo", "employees": 11}}.
Examples:
```python
# Multivariate feature flag
from posthog import get_feature_flag, get_feature_flag_payload
from hanzo_insights import get_feature_flag, get_feature_flag_payload
enabled_variant = get_feature_flag('flag-key', 'distinct_id_of_your_user')
if enabled_variant == 'variant-key':
matched_flag_payload = get_feature_flag_payload('flag-key', 'distinct_id_of_your_user')
@@ -519,6 +598,7 @@ def get_feature_flag(
only_evaluate_locally=only_evaluate_locally,
send_feature_flag_events=send_feature_flag_events,
disable_geoip=disable_geoip,
device_id=device_id,
)
@@ -529,6 +609,7 @@ def get_all_flags(
group_properties=None, # type: Optional[dict]
only_evaluate_locally=False, # type: bool
disable_geoip=None, # type: Optional[bool]
device_id=None, # type: Optional[str]
) -> Optional[dict[str, FeatureFlag]]:
"""
Get all flags for a given user.
@@ -547,7 +628,7 @@ def get_all_flags(
Examples:
```python
# All flags for user
from posthog import get_all_flags
from hanzo_insights import get_all_flags
get_all_flags('distinct_id_of_your_user')
```
Category:
@@ -561,6 +642,7 @@ def get_all_flags(
group_properties=group_properties or {},
only_evaluate_locally=only_evaluate_locally,
disable_geoip=disable_geoip,
device_id=device_id,
)
@@ -573,6 +655,7 @@ def get_feature_flag_result(
only_evaluate_locally=False,
send_feature_flag_events=True,
disable_geoip=None, # type: Optional[bool]
device_id=None, # type: Optional[str]
):
# type: (...) -> Optional[FeatureFlagResult]
"""
@@ -587,7 +670,7 @@ def get_feature_flag_result(
Example:
```python
result = posthog.get_feature_flag_result('beta-feature', 'distinct_id')
result = hanzo_insights.get_feature_flag_result('beta-feature', 'distinct_id')
if result and result.enabled:
# Use the variant and payload
print(f"Variant: {result.variant}")
@@ -604,6 +687,7 @@ def get_feature_flag_result(
only_evaluate_locally=only_evaluate_locally,
send_feature_flag_events=send_feature_flag_events,
disable_geoip=disable_geoip,
device_id=device_id,
)
@@ -617,6 +701,7 @@ def get_feature_flag_payload(
only_evaluate_locally=False,
send_feature_flag_events=True,
disable_geoip=None, # type: Optional[bool]
device_id=None, # type: Optional[str]
) -> Optional[str]:
return _proxy(
"get_feature_flag_payload",
@@ -629,6 +714,7 @@ def get_feature_flag_payload(
only_evaluate_locally=only_evaluate_locally,
send_feature_flag_events=send_feature_flag_events,
disable_geoip=disable_geoip,
device_id=device_id,
)
@@ -659,6 +745,7 @@ def get_all_flags_and_payloads(
group_properties=None, # type: Optional[dict]
only_evaluate_locally=False,
disable_geoip=None, # type: Optional[bool]
device_id=None, # type: Optional[str]
) -> FlagsAndPayloads:
return _proxy(
"get_all_flags_and_payloads",
@@ -668,6 +755,7 @@ def get_all_flags_and_payloads(
group_properties=group_properties or {},
only_evaluate_locally=only_evaluate_locally,
disable_geoip=disable_geoip,
device_id=device_id,
)
@@ -680,7 +768,7 @@ def feature_flag_definitions():
Examples:
```python
from posthog import feature_flag_definitions
from hanzo_insights import feature_flag_definitions
definitions = feature_flag_definitions()
```
@@ -692,11 +780,11 @@ def feature_flag_definitions():
def load_feature_flags():
"""
Load feature flag definitions from PostHog.
Load feature flag definitions from the server.
Examples:
```python
from posthog import load_feature_flags
from hanzo_insights import load_feature_flags
load_feature_flags()
```
@@ -712,7 +800,7 @@ def flush():
Examples:
```python
from posthog import flush
from hanzo_insights import flush
flush()
```
@@ -728,7 +816,7 @@ def join():
Examples:
```python
from posthog import join
from hanzo_insights import join
join()
```
@@ -744,7 +832,7 @@ def shutdown():
Examples:
```python
from posthog import shutdown
from hanzo_insights import shutdown
shutdown()
```
@@ -782,6 +870,7 @@ def setup() -> Client:
capture_exception_code_variables=capture_exception_code_variables,
code_variables_mask_patterns=code_variables_mask_patterns,
code_variables_ignore_patterns=code_variables_ignore_patterns,
in_app_modules=in_app_modules,
)
# always set incase user changes it
@@ -799,5 +888,7 @@ def _proxy(method, *args, **kwargs):
return fn(*args, **kwargs)
class Posthog(Client):
class Insights(Client):
"""Hanzo Insights client for product analytics."""
pass
+3
View File
@@ -0,0 +1,3 @@
from hanzo_insights.ai.prompts import Prompts
__all__ = ["Prompts"]
@@ -10,38 +10,38 @@ import time
import uuid
from typing import Any, Dict, List, Optional
from posthog.ai.types import StreamingContentBlock, TokenUsage, ToolInProgress
from posthog.ai.utils import (
from hanzo_insights.ai.types import StreamingContentBlock, TokenUsage, ToolInProgress
from hanzo_insights.ai.utils import (
call_llm_and_track_usage,
merge_usage_stats,
)
from posthog.ai.anthropic.anthropic_converter import (
from hanzo_insights.ai.anthropic.anthropic_converter import (
extract_anthropic_usage_from_event,
handle_anthropic_content_block_start,
handle_anthropic_text_delta,
handle_anthropic_tool_delta,
finalize_anthropic_tool_input,
)
from posthog.ai.sanitization import sanitize_anthropic
from posthog.client import Client as PostHogClient
from posthog import setup
from hanzo_insights.ai.sanitization import sanitize_anthropic
from hanzo_insights.client import Client as InsightsClient
from hanzo_insights import setup
class Anthropic(anthropic.Anthropic):
"""
A wrapper around the Anthropic SDK that automatically sends LLM usage events to PostHog.
A wrapper around the Anthropic SDK that automatically sends LLM usage events to Insights.
"""
_ph_client: PostHogClient
_ph_client: InsightsClient
def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs):
def __init__(self, insights_client: Optional[InsightsClient] = None, **kwargs):
"""
Args:
posthog_client: PostHog client for tracking usage
insights_client: Insights client for tracking usage
**kwargs: Additional arguments passed to the Anthropic client
"""
super().__init__(**kwargs)
self._ph_client = posthog_client or setup()
self._ph_client = insights_client or setup()
self.messages = WrappedMessages(self)
@@ -50,46 +50,46 @@ class WrappedMessages(Messages):
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,
insights_distinct_id: Optional[str] = None,
insights_trace_id: Optional[str] = None,
insights_properties: Optional[Dict[str, Any]] = None,
insights_privacy_mode: bool = False,
insights_groups: Optional[Dict[str, Any]] = None,
**kwargs: Any,
):
"""
Create a message using Anthropic's API while tracking usage in PostHog.
Create a message using Anthropic's API while tracking usage in Insights.
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
insights_distinct_id: Optional ID to associate with the usage event
insights_trace_id: Optional trace UUID for linking events
insights_properties: Optional dictionary of extra properties to include in the event
insights_privacy_mode: Whether to redact sensitive information in tracking
insights_groups: Optional group analytics properties
**kwargs: Arguments passed to Anthropic's messages.create
"""
if posthog_trace_id is None:
posthog_trace_id = str(uuid.uuid4())
if insights_trace_id is None:
insights_trace_id = str(uuid.uuid4())
if kwargs.get("stream", False):
return self._create_streaming(
posthog_distinct_id,
posthog_trace_id,
posthog_properties,
posthog_privacy_mode,
posthog_groups,
insights_distinct_id,
insights_trace_id,
insights_properties,
insights_privacy_mode,
insights_groups,
**kwargs,
)
return call_llm_and_track_usage(
posthog_distinct_id,
insights_distinct_id,
self._client._ph_client,
"anthropic",
posthog_trace_id,
posthog_properties,
posthog_privacy_mode,
posthog_groups,
insights_trace_id,
insights_properties,
insights_privacy_mode,
insights_groups,
self._client.base_url,
super().create,
**kwargs,
@@ -97,32 +97,32 @@ class WrappedMessages(Messages):
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,
insights_distinct_id: Optional[str] = None,
insights_trace_id: Optional[str] = None,
insights_properties: Optional[Dict[str, Any]] = None,
insights_privacy_mode: bool = False,
insights_groups: Optional[Dict[str, Any]] = None,
**kwargs: Any,
):
if posthog_trace_id is None:
posthog_trace_id = str(uuid.uuid4())
if insights_trace_id is None:
insights_trace_id = str(uuid.uuid4())
return self._create_streaming(
posthog_distinct_id,
posthog_trace_id,
posthog_properties,
posthog_privacy_mode,
posthog_groups,
insights_distinct_id,
insights_trace_id,
insights_properties,
insights_privacy_mode,
insights_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]],
insights_distinct_id: Optional[str],
insights_trace_id: Optional[str],
insights_properties: Optional[Dict[str, Any]],
insights_privacy_mode: bool,
insights_groups: Optional[Dict[str, Any]],
**kwargs: Any,
):
start_time = time.time()
@@ -188,11 +188,11 @@ class WrappedMessages(Messages):
latency = end_time - start_time
self._capture_streaming_event(
posthog_distinct_id,
posthog_trace_id,
posthog_properties,
posthog_privacy_mode,
posthog_groups,
insights_distinct_id,
insights_trace_id,
insights_properties,
insights_privacy_mode,
insights_groups,
kwargs,
usage_stats,
latency,
@@ -204,23 +204,23 @@ class WrappedMessages(Messages):
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]],
insights_distinct_id: Optional[str],
insights_trace_id: Optional[str],
insights_properties: Optional[Dict[str, Any]],
insights_privacy_mode: bool,
insights_groups: Optional[Dict[str, Any]],
kwargs: Dict[str, Any],
usage_stats: TokenUsage,
latency: float,
content_blocks: List[StreamingContentBlock],
accumulated_content: str,
):
from posthog.ai.types import StreamingEventData
from posthog.ai.anthropic.anthropic_converter import (
from hanzo_insights.ai.types import StreamingEventData
from hanzo_insights.ai.anthropic.anthropic_converter import (
format_anthropic_streaming_input,
format_anthropic_streaming_output_complete,
)
from posthog.ai.utils import capture_streaming_event
from hanzo_insights.ai.utils import capture_streaming_event
# Prepare standardized event data
formatted_input = format_anthropic_streaming_input(kwargs)
@@ -237,11 +237,11 @@ class WrappedMessages(Messages):
),
usage_stats=usage_stats,
latency=latency,
distinct_id=posthog_distinct_id,
trace_id=posthog_trace_id,
properties=posthog_properties,
privacy_mode=posthog_privacy_mode,
groups=posthog_groups,
distinct_id=insights_distinct_id,
trace_id=insights_trace_id,
properties=insights_properties,
privacy_mode=insights_privacy_mode,
groups=insights_groups,
)
# Use the common capture function
@@ -10,38 +10,38 @@ import time
import uuid
from typing import Any, Dict, List, Optional
from posthog import setup
from posthog.ai.types import StreamingContentBlock, TokenUsage, ToolInProgress
from posthog.ai.utils import (
from hanzo_insights import setup
from hanzo_insights.ai.types import StreamingContentBlock, TokenUsage, ToolInProgress
from hanzo_insights.ai.utils import (
call_llm_and_track_usage_async,
merge_usage_stats,
)
from posthog.ai.anthropic.anthropic_converter import (
from hanzo_insights.ai.anthropic.anthropic_converter import (
extract_anthropic_usage_from_event,
handle_anthropic_content_block_start,
handle_anthropic_text_delta,
handle_anthropic_tool_delta,
finalize_anthropic_tool_input,
)
from posthog.ai.sanitization import sanitize_anthropic
from posthog.client import Client as PostHogClient
from hanzo_insights.ai.sanitization import sanitize_anthropic
from hanzo_insights.client import Client as InsightsClient
class AsyncAnthropic(anthropic.AsyncAnthropic):
"""
An async wrapper around the Anthropic SDK that automatically sends LLM usage events to PostHog.
An async wrapper around the Anthropic SDK that automatically sends LLM usage events to Insights.
"""
_ph_client: PostHogClient
_ph_client: InsightsClient
def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs):
def __init__(self, insights_client: Optional[InsightsClient] = None, **kwargs):
"""
Args:
posthog_client: PostHog client for tracking usage
insights_client: Insights client for tracking usage
**kwargs: Additional arguments passed to the Anthropic client
"""
super().__init__(**kwargs)
self._ph_client = posthog_client or setup()
self._ph_client = insights_client or setup()
self.messages = AsyncWrappedMessages(self)
@@ -50,46 +50,46 @@ class AsyncWrappedMessages(AsyncMessages):
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,
insights_distinct_id: Optional[str] = None,
insights_trace_id: Optional[str] = None,
insights_properties: Optional[Dict[str, Any]] = None,
insights_privacy_mode: bool = False,
insights_groups: Optional[Dict[str, Any]] = None,
**kwargs: Any,
):
"""
Create a message using Anthropic's API while tracking usage in PostHog.
Create a message using Anthropic's API while tracking usage in Insights.
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
insights_distinct_id: Optional ID to associate with the usage event
insights_trace_id: Optional trace UUID for linking events
insights_properties: Optional dictionary of extra properties to include in the event
insights_privacy_mode: Whether to redact sensitive information in tracking
insights_groups: Optional group analytics properties
**kwargs: Arguments passed to Anthropic's messages.create
"""
if posthog_trace_id is None:
posthog_trace_id = str(uuid.uuid4())
if insights_trace_id is None:
insights_trace_id = str(uuid.uuid4())
if kwargs.get("stream", False):
return await self._create_streaming(
posthog_distinct_id,
posthog_trace_id,
posthog_properties,
posthog_privacy_mode,
posthog_groups,
insights_distinct_id,
insights_trace_id,
insights_properties,
insights_privacy_mode,
insights_groups,
**kwargs,
)
return await call_llm_and_track_usage_async(
posthog_distinct_id,
insights_distinct_id,
self._client._ph_client,
"anthropic",
posthog_trace_id,
posthog_properties,
posthog_privacy_mode,
posthog_groups,
insights_trace_id,
insights_properties,
insights_privacy_mode,
insights_groups,
self._client.base_url,
super().create,
**kwargs,
@@ -97,32 +97,32 @@ class AsyncWrappedMessages(AsyncMessages):
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,
insights_distinct_id: Optional[str] = None,
insights_trace_id: Optional[str] = None,
insights_properties: Optional[Dict[str, Any]] = None,
insights_privacy_mode: bool = False,
insights_groups: Optional[Dict[str, Any]] = None,
**kwargs: Any,
):
if posthog_trace_id is None:
posthog_trace_id = str(uuid.uuid4())
if insights_trace_id is None:
insights_trace_id = str(uuid.uuid4())
return await self._create_streaming(
posthog_distinct_id,
posthog_trace_id,
posthog_properties,
posthog_privacy_mode,
posthog_groups,
insights_distinct_id,
insights_trace_id,
insights_properties,
insights_privacy_mode,
insights_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]],
insights_distinct_id: Optional[str],
insights_trace_id: Optional[str],
insights_properties: Optional[Dict[str, Any]],
insights_privacy_mode: bool,
insights_groups: Optional[Dict[str, Any]],
**kwargs: Any,
):
start_time = time.time()
@@ -188,11 +188,11 @@ class AsyncWrappedMessages(AsyncMessages):
latency = end_time - start_time
await self._capture_streaming_event(
posthog_distinct_id,
posthog_trace_id,
posthog_properties,
posthog_privacy_mode,
posthog_groups,
insights_distinct_id,
insights_trace_id,
insights_properties,
insights_privacy_mode,
insights_groups,
kwargs,
usage_stats,
latency,
@@ -204,23 +204,23 @@ class AsyncWrappedMessages(AsyncMessages):
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]],
insights_distinct_id: Optional[str],
insights_trace_id: Optional[str],
insights_properties: Optional[Dict[str, Any]],
insights_privacy_mode: bool,
insights_groups: Optional[Dict[str, Any]],
kwargs: Dict[str, Any],
usage_stats: TokenUsage,
latency: float,
content_blocks: List[StreamingContentBlock],
accumulated_content: str,
):
from posthog.ai.types import StreamingEventData
from posthog.ai.anthropic.anthropic_converter import (
from hanzo_insights.ai.types import StreamingEventData
from hanzo_insights.ai.anthropic.anthropic_converter import (
format_anthropic_streaming_input,
format_anthropic_streaming_output_complete,
)
from posthog.ai.utils import capture_streaming_event
from hanzo_insights.ai.utils import capture_streaming_event
# Prepare standardized event data
formatted_input = format_anthropic_streaming_input(kwargs)
@@ -237,11 +237,11 @@ class AsyncWrappedMessages(AsyncMessages):
),
usage_stats=usage_stats,
latency=latency,
distinct_id=posthog_distinct_id,
trace_id=posthog_trace_id,
properties=posthog_properties,
privacy_mode=posthog_privacy_mode,
groups=posthog_groups,
distinct_id=insights_distinct_id,
trace_id=insights_trace_id,
properties=insights_properties,
privacy_mode=insights_privacy_mode,
groups=insights_groups,
)
# Use the common capture function
@@ -2,13 +2,13 @@
Anthropic-specific conversion utilities.
This module handles the conversion of Anthropic API responses and inputs
into standardized formats for PostHog tracking.
into standardized formats for Insights tracking.
"""
import json
from typing import Any, Dict, List, Optional, Tuple
from posthog.ai.types import (
from hanzo_insights.ai.types import (
FormattedContentItem,
FormattedFunctionCall,
FormattedMessage,
@@ -17,6 +17,7 @@ from posthog.ai.types import (
TokenUsage,
ToolInProgress,
)
from hanzo_insights.ai.utils import serialize_raw_usage
def format_anthropic_response(response: Any) -> List[FormattedMessage]:
@@ -221,6 +222,12 @@ def extract_anthropic_usage_from_response(response: Any) -> TokenUsage:
if web_search_count > 0:
result["web_search_count"] = web_search_count
# Capture raw usage metadata for backend processing
# Serialize to dict here in the converter (not in utils)
serialized = serialize_raw_usage(response.usage)
if serialized:
result["raw_usage"] = serialized
return result
@@ -247,6 +254,11 @@ def extract_anthropic_usage_from_event(event: Any) -> TokenUsage:
usage["cache_read_input_tokens"] = getattr(
event.message.usage, "cache_read_input_tokens", 0
)
# Capture raw usage metadata for backend processing
# Serialize to dict here in the converter (not in utils)
serialized = serialize_raw_usage(event.message.usage)
if serialized:
usage["raw_usage"] = serialized
# Handle usage stats from message_delta event
if hasattr(event, "usage") and event.usage:
@@ -262,6 +274,12 @@ def extract_anthropic_usage_from_event(event: Any) -> TokenUsage:
if web_search_count > 0:
usage["web_search_count"] = web_search_count
# Capture raw usage metadata for backend processing
# Serialize to dict here in the converter (not in utils)
serialized = serialize_raw_usage(event.usage)
if serialized:
usage["raw_usage"] = serialized
return usage
@@ -407,9 +425,9 @@ def format_anthropic_streaming_input(kwargs: Dict[str, Any]) -> Any:
kwargs: Keyword arguments passed to Anthropic API
Returns:
Formatted input ready for PostHog tracking
Formatted input ready for Insights tracking
"""
from posthog.ai.utils import merge_system_prompt
from hanzo_insights.ai.utils import merge_system_prompt
return merge_system_prompt(kwargs, "anthropic")
@@ -427,7 +445,7 @@ def format_anthropic_streaming_output_complete(
accumulated_content: Raw accumulated text content as fallback
Returns:
Formatted messages ready for PostHog tracking
Formatted messages ready for Insights tracking
"""
formatted_content = format_anthropic_streaming_content(content_blocks)
@@ -7,59 +7,59 @@ except ImportError:
from typing import Optional
from posthog.ai.anthropic.anthropic import WrappedMessages
from posthog.ai.anthropic.anthropic_async import AsyncWrappedMessages
from posthog.client import Client as PostHogClient
from posthog import setup
from hanzo_insights.ai.anthropic.anthropic import WrappedMessages
from hanzo_insights.ai.anthropic.anthropic_async import AsyncWrappedMessages
from hanzo_insights.client import Client as InsightsClient
from hanzo_insights import setup
class AnthropicBedrock(anthropic.AnthropicBedrock):
"""
A wrapper around the Anthropic Bedrock SDK that automatically sends LLM usage events to PostHog.
A wrapper around the Anthropic Bedrock SDK that automatically sends LLM usage events to Insights.
"""
_ph_client: PostHogClient
_ph_client: InsightsClient
def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs):
def __init__(self, insights_client: Optional[InsightsClient] = None, **kwargs):
super().__init__(**kwargs)
self._ph_client = posthog_client or setup()
self._ph_client = insights_client or setup()
self.messages = WrappedMessages(self)
class AsyncAnthropicBedrock(anthropic.AsyncAnthropicBedrock):
"""
A wrapper around the Anthropic Bedrock SDK that automatically sends LLM usage events to PostHog.
A wrapper around the Anthropic Bedrock SDK that automatically sends LLM usage events to Insights.
"""
_ph_client: PostHogClient
_ph_client: InsightsClient
def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs):
def __init__(self, insights_client: Optional[InsightsClient] = None, **kwargs):
super().__init__(**kwargs)
self._ph_client = posthog_client or setup()
self._ph_client = insights_client or setup()
self.messages = AsyncWrappedMessages(self)
class AnthropicVertex(anthropic.AnthropicVertex):
"""
A wrapper around the Anthropic Vertex SDK that automatically sends LLM usage events to PostHog.
A wrapper around the Anthropic Vertex SDK that automatically sends LLM usage events to Insights.
"""
_ph_client: PostHogClient
_ph_client: InsightsClient
def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs):
def __init__(self, insights_client: Optional[InsightsClient] = None, **kwargs):
super().__init__(**kwargs)
self._ph_client = posthog_client or setup()
self._ph_client = insights_client or setup()
self.messages = WrappedMessages(self)
class AsyncAnthropicVertex(anthropic.AsyncAnthropicVertex):
"""
A wrapper around the Anthropic Vertex SDK that automatically sends LLM usage events to PostHog.
A wrapper around the Anthropic Vertex SDK that automatically sends LLM usage events to Insights.
"""
_ph_client: PostHogClient
_ph_client: InsightsClient
def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs):
def __init__(self, insights_client: Optional[InsightsClient] = None, **kwargs):
super().__init__(**kwargs)
self._ph_client = posthog_client or setup()
self._ph_client = insights_client or setup()
self.messages = AsyncWrappedMessages(self)
@@ -1,4 +1,5 @@
from .gemini import Client
from .gemini_async import AsyncClient
from .gemini_converter import (
format_gemini_input,
format_gemini_response,
@@ -9,12 +10,14 @@ from .gemini_converter import (
# Create a genai-like module for perfect drop-in replacement
class _GenAI:
Client = Client
AsyncClient = AsyncClient
genai = _GenAI()
__all__ = [
"Client",
"AsyncClient",
"genai",
"format_gemini_input",
"format_gemini_response",
@@ -3,8 +3,8 @@ import time
import uuid
from typing import Any, Dict, Optional
from posthog.ai.types import TokenUsage, StreamingEventData
from posthog.ai.utils import merge_system_prompt
from hanzo_insights.ai.types import TokenUsage, StreamingEventData
from hanzo_insights.ai.utils import merge_system_prompt
try:
from google import genai
@@ -13,40 +13,40 @@ except ImportError:
"Please install the Google Gemini SDK to use this feature: 'pip install google-genai'"
)
from posthog import setup
from posthog.ai.utils import (
from hanzo_insights import setup
from hanzo_insights.ai.utils import (
call_llm_and_track_usage,
capture_streaming_event,
merge_usage_stats,
)
from posthog.ai.gemini.gemini_converter import (
from hanzo_insights.ai.gemini.gemini_converter import (
extract_gemini_usage_from_chunk,
extract_gemini_content_from_chunk,
format_gemini_streaming_output,
)
from posthog.ai.sanitization import sanitize_gemini
from posthog.client import Client as PostHogClient
from hanzo_insights.ai.sanitization import sanitize_gemini
from hanzo_insights.client import Client as InsightsClient
class Client:
"""
A drop-in replacement for genai.Client that automatically sends LLM usage events to PostHog.
A drop-in replacement for genai.Client that automatically sends LLM usage events to Insights.
Usage:
client = Client(
api_key="your_api_key",
posthog_client=posthog_client,
posthog_distinct_id="default_user", # Optional defaults
posthog_properties={"team": "ai"} # Optional defaults
insights_client=insights_client,
insights_distinct_id="default_user", # Optional defaults
insights_properties={"team": "ai"} # Optional defaults
)
response = client.models.generate_content(
model="gemini-2.0-flash",
contents=["Hello world"],
posthog_distinct_id="specific_user" # Override default
insights_distinct_id="specific_user" # Override default
)
"""
_ph_client: PostHogClient
_ph_client: InsightsClient
def __init__(
self,
@@ -57,11 +57,11 @@ class Client:
location: Optional[str] = None,
debug_config: Optional[Any] = None,
http_options: Optional[Any] = None,
posthog_client: Optional[PostHogClient] = None,
posthog_distinct_id: Optional[str] = None,
posthog_properties: Optional[Dict[str, Any]] = None,
posthog_privacy_mode: bool = False,
posthog_groups: Optional[Dict[str, Any]] = None,
insights_client: Optional[InsightsClient] = None,
insights_distinct_id: Optional[str] = None,
insights_properties: Optional[Dict[str, Any]] = None,
insights_privacy_mode: bool = False,
insights_groups: Optional[Dict[str, Any]] = None,
**kwargs,
):
"""
@@ -73,18 +73,18 @@ class Client:
location: GCP location for Vertex AI
debug_config: Debug configuration for the client
http_options: HTTP options for the client
posthog_client: PostHog client for tracking usage
posthog_distinct_id: Default distinct ID for all calls (can be overridden per call)
posthog_properties: Default properties for all calls (can be overridden per call)
posthog_privacy_mode: Default privacy mode for all calls (can be overridden per call)
posthog_groups: Default groups for all calls (can be overridden per call)
insights_client: Insights client for tracking usage
insights_distinct_id: Default distinct ID for all calls (can be overridden per call)
insights_properties: Default properties for all calls (can be overridden per call)
insights_privacy_mode: Default privacy mode for all calls (can be overridden per call)
insights_groups: Default groups for all calls (can be overridden per call)
**kwargs: Additional arguments (for future compatibility)
"""
self._ph_client = posthog_client or setup()
self._ph_client = insights_client or setup()
if self._ph_client is None:
raise ValueError("posthog_client is required for PostHog tracking")
raise ValueError("insights_client is required for Insights tracking")
self.models = Models(
api_key=api_key,
@@ -94,21 +94,21 @@ class Client:
location=location,
debug_config=debug_config,
http_options=http_options,
posthog_client=self._ph_client,
posthog_distinct_id=posthog_distinct_id,
posthog_properties=posthog_properties,
posthog_privacy_mode=posthog_privacy_mode,
posthog_groups=posthog_groups,
insights_client=self._ph_client,
insights_distinct_id=insights_distinct_id,
insights_properties=insights_properties,
insights_privacy_mode=insights_privacy_mode,
insights_groups=insights_groups,
**kwargs,
)
class Models:
"""
Models interface that mimics genai.Client().models with PostHog tracking.
Models interface that mimics genai.Client().models with Insights tracking.
"""
_ph_client: PostHogClient # Not None after __init__ validation
_ph_client: InsightsClient # Not None after __init__ validation
def __init__(
self,
@@ -119,11 +119,11 @@ class Models:
location: Optional[str] = None,
debug_config: Optional[Any] = None,
http_options: Optional[Any] = None,
posthog_client: Optional[PostHogClient] = None,
posthog_distinct_id: Optional[str] = None,
posthog_properties: Optional[Dict[str, Any]] = None,
posthog_privacy_mode: bool = False,
posthog_groups: Optional[Dict[str, Any]] = None,
insights_client: Optional[InsightsClient] = None,
insights_distinct_id: Optional[str] = None,
insights_properties: Optional[Dict[str, Any]] = None,
insights_privacy_mode: bool = False,
insights_groups: Optional[Dict[str, Any]] = None,
**kwargs,
):
"""
@@ -135,24 +135,24 @@ class Models:
location: GCP location for Vertex AI
debug_config: Debug configuration for the client
http_options: HTTP options for the client
posthog_client: PostHog client for tracking usage
posthog_distinct_id: Default distinct ID for all calls
posthog_properties: Default properties for all calls
posthog_privacy_mode: Default privacy mode for all calls
posthog_groups: Default groups for all calls
insights_client: Insights client for tracking usage
insights_distinct_id: Default distinct ID for all calls
insights_properties: Default properties for all calls
insights_privacy_mode: Default privacy mode for all calls
insights_groups: Default groups for all calls
**kwargs: Additional arguments (for future compatibility)
"""
self._ph_client = posthog_client or setup()
self._ph_client = insights_client or setup()
if self._ph_client is None:
raise ValueError("posthog_client is required for PostHog tracking")
raise ValueError("insights_client is required for Insights tracking")
# Store default PostHog settings
self._default_distinct_id = posthog_distinct_id
self._default_properties = posthog_properties or {}
self._default_privacy_mode = posthog_privacy_mode
self._default_groups = posthog_groups
# Store default Insights settings
self._default_distinct_id = insights_distinct_id
self._default_properties = insights_properties or {}
self._default_privacy_mode = insights_privacy_mode
self._default_groups = insights_groups
# Build genai.Client arguments
client_args: Dict[str, Any] = {}
@@ -196,7 +196,7 @@ class Models:
self._client = genai.Client(**client_args)
self._base_url = "https://generativelanguage.googleapis.com"
def _merge_posthog_params(
def _merge_insights_params(
self,
call_distinct_id: Optional[str],
call_trace_id: Optional[str],
@@ -204,7 +204,7 @@ class Models:
call_privacy_mode: Optional[bool],
call_groups: Optional[Dict[str, Any]],
):
"""Merge call-level PostHog parameters with client defaults."""
"""Merge call-level Insights parameters with client defaults."""
# Use call-level values if provided, otherwise fall back to defaults
distinct_id = (
@@ -234,38 +234,38 @@ class Models:
self,
model: str,
contents,
posthog_distinct_id: Optional[str] = None,
posthog_trace_id: Optional[str] = None,
posthog_properties: Optional[Dict[str, Any]] = None,
posthog_privacy_mode: Optional[bool] = None,
posthog_groups: Optional[Dict[str, Any]] = None,
insights_distinct_id: Optional[str] = None,
insights_trace_id: Optional[str] = None,
insights_properties: Optional[Dict[str, Any]] = None,
insights_privacy_mode: Optional[bool] = None,
insights_groups: Optional[Dict[str, Any]] = None,
**kwargs: Any,
):
"""
Generate content using Gemini's API while tracking usage in PostHog.
Generate content using Gemini's API while tracking usage in Insights.
This method signature exactly matches genai.Client().models.generate_content()
with additional PostHog tracking parameters.
with additional Insights tracking parameters.
Args:
model: The model to use (e.g., 'gemini-2.0-flash')
contents: The input content for generation
posthog_distinct_id: ID to associate with the usage event (overrides client default)
posthog_trace_id: Trace UUID for linking events (auto-generated if not provided)
posthog_properties: Extra properties to include in the event (merged with client defaults)
posthog_privacy_mode: Whether to redact sensitive information (overrides client default)
posthog_groups: Group analytics properties (overrides client default)
insights_distinct_id: ID to associate with the usage event (overrides client default)
insights_trace_id: Trace UUID for linking events (auto-generated if not provided)
insights_properties: Extra properties to include in the event (merged with client defaults)
insights_privacy_mode: Whether to redact sensitive information (overrides client default)
insights_groups: Group analytics properties (overrides client default)
**kwargs: Arguments passed to Gemini's generate_content
"""
# Merge PostHog parameters
# Merge Insights parameters
distinct_id, trace_id, properties, privacy_mode, groups = (
self._merge_posthog_params(
posthog_distinct_id,
posthog_trace_id,
posthog_properties,
posthog_privacy_mode,
posthog_groups,
self._merge_insights_params(
insights_distinct_id,
insights_trace_id,
insights_properties,
insights_privacy_mode,
insights_groups,
)
)
@@ -304,7 +304,7 @@ class Models:
def generator():
nonlocal usage_stats
nonlocal accumulated_content # noqa: F824
nonlocal accumulated_content
try:
for chunk in response:
# Extract usage stats from chunk
@@ -380,7 +380,7 @@ class Models:
capture_streaming_event(self._ph_client, event_data)
def _format_input(self, contents, **kwargs):
"""Format input contents for PostHog tracking"""
"""Format input contents for Insights tracking"""
# Create kwargs dict with contents for merge_system_prompt
input_kwargs = {"contents": contents, **kwargs}
@@ -390,21 +390,21 @@ class Models:
self,
model: str,
contents,
posthog_distinct_id: Optional[str] = None,
posthog_trace_id: Optional[str] = None,
posthog_properties: Optional[Dict[str, Any]] = None,
posthog_privacy_mode: Optional[bool] = None,
posthog_groups: Optional[Dict[str, Any]] = None,
insights_distinct_id: Optional[str] = None,
insights_trace_id: Optional[str] = None,
insights_properties: Optional[Dict[str, Any]] = None,
insights_privacy_mode: Optional[bool] = None,
insights_groups: Optional[Dict[str, Any]] = None,
**kwargs: Any,
):
# Merge PostHog parameters
# Merge Insights parameters
distinct_id, trace_id, properties, privacy_mode, groups = (
self._merge_posthog_params(
posthog_distinct_id,
posthog_trace_id,
posthog_properties,
posthog_privacy_mode,
posthog_groups,
self._merge_insights_params(
insights_distinct_id,
insights_trace_id,
insights_properties,
insights_privacy_mode,
insights_groups,
)
)
+423
View File
@@ -0,0 +1,423 @@
import os
import time
import uuid
from typing import Any, Dict, Optional
from hanzo_insights.ai.types import TokenUsage, StreamingEventData
from hanzo_insights.ai.utils import merge_system_prompt
try:
from google import genai
except ImportError:
raise ModuleNotFoundError(
"Please install the Google Gemini SDK to use this feature: 'pip install google-genai'"
)
from hanzo_insights import setup
from hanzo_insights.ai.utils import (
call_llm_and_track_usage_async,
capture_streaming_event,
merge_usage_stats,
)
from hanzo_insights.ai.gemini.gemini_converter import (
extract_gemini_usage_from_chunk,
extract_gemini_content_from_chunk,
format_gemini_streaming_output,
)
from hanzo_insights.ai.sanitization import sanitize_gemini
from hanzo_insights.client import Client as InsightsClient
class AsyncClient:
"""
An async drop-in replacement for genai.Client that automatically sends LLM usage events to Insights.
Usage:
client = AsyncClient(
api_key="your_api_key",
insights_client=insights_client,
insights_distinct_id="default_user", # Optional defaults
insights_properties={"team": "ai"} # Optional defaults
)
response = await client.models.generate_content(
model="gemini-2.0-flash",
contents=["Hello world"],
insights_distinct_id="specific_user" # Override default
)
"""
_ph_client: InsightsClient
def __init__(
self,
api_key: Optional[str] = None,
vertexai: Optional[bool] = None,
credentials: Optional[Any] = None,
project: Optional[str] = None,
location: Optional[str] = None,
debug_config: Optional[Any] = None,
http_options: Optional[Any] = None,
insights_client: Optional[InsightsClient] = None,
insights_distinct_id: Optional[str] = None,
insights_properties: Optional[Dict[str, Any]] = None,
insights_privacy_mode: bool = False,
insights_groups: Optional[Dict[str, Any]] = None,
**kwargs,
):
"""
Args:
api_key: Google AI API key. If not provided, will use GOOGLE_API_KEY or API_KEY environment variable (not required for Vertex AI)
vertexai: Whether to use Vertex AI authentication
credentials: Vertex AI credentials object
project: GCP project ID for Vertex AI
location: GCP location for Vertex AI
debug_config: Debug configuration for the client
http_options: HTTP options for the client
insights_client: Insights client for tracking usage
insights_distinct_id: Default distinct ID for all calls (can be overridden per call)
insights_properties: Default properties for all calls (can be overridden per call)
insights_privacy_mode: Default privacy mode for all calls (can be overridden per call)
insights_groups: Default groups for all calls (can be overridden per call)
**kwargs: Additional arguments (for future compatibility)
"""
self._ph_client = insights_client or setup()
if self._ph_client is None:
raise ValueError("insights_client is required for Insights tracking")
self.models = AsyncModels(
api_key=api_key,
vertexai=vertexai,
credentials=credentials,
project=project,
location=location,
debug_config=debug_config,
http_options=http_options,
insights_client=self._ph_client,
insights_distinct_id=insights_distinct_id,
insights_properties=insights_properties,
insights_privacy_mode=insights_privacy_mode,
insights_groups=insights_groups,
**kwargs,
)
class AsyncModels:
"""
Async Models interface that mimics genai.Client().aio.models with Insights tracking.
"""
_ph_client: InsightsClient # Not None after __init__ validation
def __init__(
self,
api_key: Optional[str] = None,
vertexai: Optional[bool] = None,
credentials: Optional[Any] = None,
project: Optional[str] = None,
location: Optional[str] = None,
debug_config: Optional[Any] = None,
http_options: Optional[Any] = None,
insights_client: Optional[InsightsClient] = None,
insights_distinct_id: Optional[str] = None,
insights_properties: Optional[Dict[str, Any]] = None,
insights_privacy_mode: bool = False,
insights_groups: Optional[Dict[str, Any]] = None,
**kwargs,
):
"""
Args:
api_key: Google AI API key. If not provided, will use GOOGLE_API_KEY or API_KEY environment variable (not required for Vertex AI)
vertexai: Whether to use Vertex AI authentication
credentials: Vertex AI credentials object
project: GCP project ID for Vertex AI
location: GCP location for Vertex AI
debug_config: Debug configuration for the client
http_options: HTTP options for the client
insights_client: Insights client for tracking usage
insights_distinct_id: Default distinct ID for all calls
insights_properties: Default properties for all calls
insights_privacy_mode: Default privacy mode for all calls
insights_groups: Default groups for all calls
**kwargs: Additional arguments (for future compatibility)
"""
self._ph_client = insights_client or setup()
if self._ph_client is None:
raise ValueError("insights_client is required for Insights tracking")
# Store default Insights settings
self._default_distinct_id = insights_distinct_id
self._default_properties = insights_properties or {}
self._default_privacy_mode = insights_privacy_mode
self._default_groups = insights_groups
# Build genai.Client arguments
client_args: Dict[str, Any] = {}
# Add Vertex AI parameters if provided
if vertexai is not None:
client_args["vertexai"] = vertexai
if credentials is not None:
client_args["credentials"] = credentials
if project is not None:
client_args["project"] = project
if location is not None:
client_args["location"] = location
if debug_config is not None:
client_args["debug_config"] = debug_config
if http_options is not None:
client_args["http_options"] = http_options
# Handle API key authentication
if vertexai:
# For Vertex AI, api_key is optional
if api_key is not None:
client_args["api_key"] = api_key
else:
# For non-Vertex AI mode, api_key is required (backwards compatibility)
if api_key is None:
api_key = os.environ.get("GOOGLE_API_KEY") or os.environ.get("API_KEY")
if api_key is None:
raise ValueError(
"API key must be provided either as parameter or via GOOGLE_API_KEY/API_KEY environment variable"
)
client_args["api_key"] = api_key
self._client = genai.Client(**client_args)
self._base_url = "https://generativelanguage.googleapis.com"
def _merge_insights_params(
self,
call_distinct_id: Optional[str],
call_trace_id: Optional[str],
call_properties: Optional[Dict[str, Any]],
call_privacy_mode: Optional[bool],
call_groups: Optional[Dict[str, Any]],
):
"""Merge call-level Insights parameters with client defaults."""
# Use call-level values if provided, otherwise fall back to defaults
distinct_id = (
call_distinct_id
if call_distinct_id is not None
else self._default_distinct_id
)
privacy_mode = (
call_privacy_mode
if call_privacy_mode is not None
else self._default_privacy_mode
)
groups = call_groups if call_groups is not None else self._default_groups
# Merge properties: default properties + call properties (call properties override)
properties = dict(self._default_properties)
if call_properties:
properties.update(call_properties)
if call_trace_id is None:
call_trace_id = str(uuid.uuid4())
return distinct_id, call_trace_id, properties, privacy_mode, groups
async def generate_content(
self,
model: str,
contents,
insights_distinct_id: Optional[str] = None,
insights_trace_id: Optional[str] = None,
insights_properties: Optional[Dict[str, Any]] = None,
insights_privacy_mode: Optional[bool] = None,
insights_groups: Optional[Dict[str, Any]] = None,
**kwargs: Any,
):
"""
Generate content using Gemini's API while tracking usage in Insights.
This method signature exactly matches genai.Client().aio.models.generate_content()
with additional Insights tracking parameters.
Args:
model: The model to use (e.g., 'gemini-2.0-flash')
contents: The input content for generation
insights_distinct_id: ID to associate with the usage event (overrides client default)
insights_trace_id: Trace UUID for linking events (auto-generated if not provided)
insights_properties: Extra properties to include in the event (merged with client defaults)
insights_privacy_mode: Whether to redact sensitive information (overrides client default)
insights_groups: Group analytics properties (overrides client default)
**kwargs: Arguments passed to Gemini's generate_content
"""
# Merge Insights parameters
distinct_id, trace_id, properties, privacy_mode, groups = (
self._merge_insights_params(
insights_distinct_id,
insights_trace_id,
insights_properties,
insights_privacy_mode,
insights_groups,
)
)
kwargs_with_contents = {"model": model, "contents": contents, **kwargs}
return await call_llm_and_track_usage_async(
distinct_id,
self._ph_client,
"gemini",
trace_id,
properties,
privacy_mode,
groups,
self._base_url,
self._client.aio.models.generate_content,
**kwargs_with_contents,
)
async def _generate_content_streaming(
self,
model: str,
contents,
distinct_id: Optional[str],
trace_id: Optional[str],
properties: Optional[Dict[str, Any]],
privacy_mode: bool,
groups: Optional[Dict[str, Any]],
**kwargs: Any,
):
start_time = time.time()
usage_stats: TokenUsage = TokenUsage(input_tokens=0, output_tokens=0)
accumulated_content = []
kwargs_without_stream = {"model": model, "contents": contents, **kwargs}
response = await self._client.aio.models.generate_content_stream(
**kwargs_without_stream
)
async def async_generator():
nonlocal usage_stats
nonlocal accumulated_content
try:
async for chunk in response:
# Extract usage stats from chunk
chunk_usage = extract_gemini_usage_from_chunk(chunk)
if chunk_usage:
# Gemini reports cumulative totals, not incremental values
merge_usage_stats(usage_stats, chunk_usage, mode="cumulative")
# Extract content from chunk (now returns content blocks)
content_block = extract_gemini_content_from_chunk(chunk)
if content_block is not None:
accumulated_content.append(content_block)
yield chunk
finally:
end_time = time.time()
latency = end_time - start_time
self._capture_streaming_event(
model,
contents,
distinct_id,
trace_id,
properties,
privacy_mode,
groups,
kwargs,
usage_stats,
latency,
accumulated_content,
)
return async_generator()
def _capture_streaming_event(
self,
model: str,
contents,
distinct_id: Optional[str],
trace_id: Optional[str],
properties: Optional[Dict[str, Any]],
privacy_mode: bool,
groups: Optional[Dict[str, Any]],
kwargs: Dict[str, Any],
usage_stats: TokenUsage,
latency: float,
output: Any,
):
# Prepare standardized event data
formatted_input = self._format_input(contents, **kwargs)
sanitized_input = sanitize_gemini(formatted_input)
event_data = StreamingEventData(
provider="gemini",
model=model,
base_url=self._base_url,
kwargs=kwargs,
formatted_input=sanitized_input,
formatted_output=format_gemini_streaming_output(output),
usage_stats=usage_stats,
latency=latency,
distinct_id=distinct_id,
trace_id=trace_id,
properties=properties,
privacy_mode=privacy_mode,
groups=groups,
)
# Use the common capture function
capture_streaming_event(self._ph_client, event_data)
def _format_input(self, contents, **kwargs):
"""Format input contents for Insights tracking"""
# Create kwargs dict with contents for merge_system_prompt
input_kwargs = {"contents": contents, **kwargs}
return merge_system_prompt(input_kwargs, "gemini")
async def generate_content_stream(
self,
model: str,
contents,
insights_distinct_id: Optional[str] = None,
insights_trace_id: Optional[str] = None,
insights_properties: Optional[Dict[str, Any]] = None,
insights_privacy_mode: Optional[bool] = None,
insights_groups: Optional[Dict[str, Any]] = None,
**kwargs: Any,
):
# Merge Insights parameters
distinct_id, trace_id, properties, privacy_mode, groups = (
self._merge_insights_params(
insights_distinct_id,
insights_trace_id,
insights_properties,
insights_privacy_mode,
insights_groups,
)
)
return await self._generate_content_streaming(
model,
contents,
distinct_id,
trace_id,
properties,
privacy_mode,
groups,
**kwargs,
)
@@ -2,16 +2,17 @@
Gemini-specific conversion utilities.
This module handles the conversion of Gemini API responses and inputs
into standardized formats for PostHog tracking.
into standardized formats for Insights tracking.
"""
from typing import Any, Dict, List, Optional, TypedDict, Union
from posthog.ai.types import (
from hanzo_insights.ai.types import (
FormattedContentItem,
FormattedMessage,
TokenUsage,
)
from hanzo_insights.ai.utils import serialize_raw_usage
class GeminiPart(TypedDict, total=False):
@@ -29,35 +30,76 @@ class GeminiMessage(TypedDict, total=False):
text: str
def _extract_text_from_parts(parts: List[Any]) -> str:
def _format_parts_as_content_blocks(parts: List[Any]) -> List[FormattedContentItem]:
"""
Extract and concatenate text from a parts array.
Format Gemini parts array into structured content blocks.
Preserves structure for multimodal content (text + images) instead of
concatenating everything into a string.
Args:
parts: List of parts that may contain text content
parts: List of parts that may contain text, inline_data, etc.
Returns:
Concatenated text from all parts
List of formatted content blocks
"""
content_parts = []
content_blocks: List[FormattedContentItem] = []
for part in parts:
# Handle dict with text field
if isinstance(part, dict) and "text" in part:
content_parts.append(part["text"])
content_blocks.append({"type": "text", "text": part["text"]})
# Handle string parts
elif isinstance(part, str):
content_parts.append(part)
content_blocks.append({"type": "text", "text": part})
# Handle dict with inline_data (images, documents, etc.)
elif isinstance(part, dict) and "inline_data" in part:
inline_data = part["inline_data"]
mime_type = inline_data.get("mime_type", "")
content_type = "image" if mime_type.startswith("image/") else "document"
content_blocks.append(
{
"type": content_type,
"inline_data": inline_data,
}
)
# Handle object with text attribute
elif hasattr(part, "text"):
# Get the text attribute value
text_value = getattr(part, "text", "")
content_parts.append(text_value if text_value else str(part))
if text_value:
content_blocks.append({"type": "text", "text": text_value})
else:
content_parts.append(str(part))
# Handle object with inline_data attribute
elif hasattr(part, "inline_data"):
inline_data = part.inline_data
# Convert to dict if needed
if hasattr(inline_data, "mime_type") and hasattr(inline_data, "data"):
# Determine type based on mime_type
mime_type = inline_data.mime_type
content_type = "image" if mime_type.startswith("image/") else "document"
return "".join(content_parts)
content_blocks.append(
{
"type": content_type,
"inline_data": {
"mime_type": mime_type,
"data": inline_data.data,
},
}
)
else:
content_blocks.append(
{
"type": "image",
"inline_data": inline_data,
}
)
return content_blocks
def _format_dict_message(item: Dict[str, Any]) -> FormattedMessage:
@@ -73,16 +115,17 @@ def _format_dict_message(item: Dict[str, Any]) -> FormattedMessage:
# Handle dict format with parts array (Gemini-specific format)
if "parts" in item and isinstance(item["parts"], list):
content = _extract_text_from_parts(item["parts"])
return {"role": item.get("role", "user"), "content": content}
content_blocks = _format_parts_as_content_blocks(item["parts"])
return {"role": item.get("role", "user"), "content": content_blocks}
# Handle dict with content field
if "content" in item:
content = item["content"]
if isinstance(content, list):
# If content is a list, extract text from it
content = _extract_text_from_parts(content)
# If content is a list, format it as content blocks
content_blocks = _format_parts_as_content_blocks(content)
return {"role": item.get("role", "user"), "content": content_blocks}
elif not isinstance(content, str):
content = str(content)
@@ -110,14 +153,14 @@ def _format_object_message(item: Any) -> FormattedMessage:
# Handle object with parts attribute
if hasattr(item, "parts") and hasattr(item.parts, "__iter__"):
content = _extract_text_from_parts(item.parts)
content_blocks = _format_parts_as_content_blocks(list(item.parts))
role = getattr(item, "role", "user") if hasattr(item, "role") else "user"
# Ensure role is a string
if not isinstance(role, str):
role = "user"
return {"role": role, "content": content}
return {"role": role, "content": content_blocks}
# Handle object with text attribute
if hasattr(item, "text"):
@@ -140,7 +183,8 @@ def _format_object_message(item: Any) -> FormattedMessage:
content = item.content
if isinstance(content, list):
content = _extract_text_from_parts(content)
content_blocks = _format_parts_as_content_blocks(content)
return {"role": role, "content": content_blocks}
elif not isinstance(content, str):
content = str(content)
@@ -193,6 +237,29 @@ def format_gemini_response(response: Any) -> List[FormattedMessage]:
}
)
elif hasattr(part, "inline_data") and part.inline_data:
# Handle audio/media inline data
import base64
inline_data = part.inline_data
mime_type = getattr(inline_data, "mime_type", "audio/pcm")
raw_data = getattr(inline_data, "data", b"")
# Encode binary data as base64 string for JSON serialization
if isinstance(raw_data, bytes):
data = base64.b64encode(raw_data).decode("utf-8")
else:
# Already a string (base64)
data = raw_data
content.append(
{
"type": "audio",
"mime_type": mime_type,
"data": data,
}
)
if content:
output.append(
{
@@ -282,7 +349,7 @@ def format_gemini_input_with_system(
if system_instruction is not None:
has_system = any(msg.get("role") == "system" for msg in formatted_messages)
if not has_system:
from posthog.ai.types import FormattedMessage
from hanzo_insights.ai.types import FormattedMessage
system_message: FormattedMessage = {
"role": "system",
@@ -295,7 +362,7 @@ def format_gemini_input_with_system(
def format_gemini_input(contents: Any) -> List[FormattedMessage]:
"""
Format Gemini input contents into standardized message format for PostHog tracking.
Format Gemini input contents into standardized message format for Insights tracking.
This function handles various input formats:
- String inputs
@@ -421,6 +488,12 @@ def _extract_usage_from_metadata(metadata: Any) -> TokenUsage:
if reasoning_tokens and reasoning_tokens > 0:
usage["reasoning_tokens"] = reasoning_tokens
# Capture raw usage metadata for backend processing
# Serialize to dict here in the converter (not in utils)
serialized = serialize_raw_usage(metadata)
if serialized:
usage["raw_usage"] = serialized
return usage
@@ -1,8 +1,8 @@
try:
import langchain # noqa: F401
import langchain_core # noqa: F401
except ImportError:
raise ModuleNotFoundError(
"Please install LangChain to use this feature: 'pip install langchain'"
"Please install LangChain to use this feature: 'pip install langchain-core'"
)
import json
@@ -22,8 +22,8 @@ from uuid import UUID
try:
# LangChain 1.0+ and modern 0.x with langchain-core
from langchain_core.callbacks.base import BaseCallbackHandler
from langchain_core.agents import AgentAction, AgentFinish
from langchain_core.callbacks.base import BaseCallbackHandler
except (ImportError, ModuleNotFoundError):
# Fallback for older LangChain versions
from langchain.callbacks.base import BaseCallbackHandler
@@ -35,18 +35,18 @@ from langchain_core.messages import (
FunctionMessage,
HumanMessage,
SystemMessage,
ToolMessage,
ToolCall,
ToolMessage,
)
from langchain_core.outputs import ChatGeneration, LLMResult
from pydantic import BaseModel
from posthog import setup
from posthog.ai.utils import get_model_params, with_privacy_mode
from posthog.ai.sanitization import sanitize_langchain
from posthog.client import Client
from hanzo_insights import setup
from hanzo_insights.ai.sanitization import sanitize_langchain
from hanzo_insights.ai.utils import get_model_params, with_privacy_mode
from hanzo_insights.client import Client
log = logging.getLogger("posthog")
log = logging.getLogger("hanzo_insights")
@dataclass
@@ -79,6 +79,8 @@ class GenerationMetadata(SpanMetadata):
"""Base URL of the provider's API used in the run."""
tools: Optional[List[Dict[str, Any]]] = None
"""Tools provided to the model."""
insights_properties: Optional[Dict[str, Any]] = None
"""Insights properties of the run."""
RunMetadata = Union[SpanMetadata, GenerationMetadata]
@@ -87,11 +89,11 @@ RunMetadataStorage = Dict[UUID, RunMetadata]
class CallbackHandler(BaseCallbackHandler):
"""
The PostHog LLM observability callback handler for LangChain.
The Insights LLM observability callback handler for LangChain.
"""
_ph_client: Client
"""PostHog client instance."""
"""Insights client instance."""
_distinct_id: Optional[Union[str, int, UUID]]
"""Distinct ID of the user to associate the trace with."""
@@ -129,12 +131,12 @@ class CallbackHandler(BaseCallbackHandler):
):
"""
Args:
client: PostHog client instance.
client: Insights 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.
groups: Optional additional Insights groups to use for the trace.
"""
self._ph_client = client or setup()
self._distinct_id = distinct_id
@@ -420,6 +422,8 @@ class CallbackHandler(BaseCallbackHandler):
generation.model = model
if provider := metadata.get("ls_provider"):
generation.provider = provider
generation.insights_properties = metadata.get("insights_properties")
try:
base_url = serialized["kwargs"]["openai_api_base"]
if base_url is not None:
@@ -502,6 +506,14 @@ class CallbackHandler(BaseCallbackHandler):
if isinstance(outputs, BaseException):
event_properties["$ai_error"] = _stringify_exception(outputs)
event_properties["$ai_is_error"] = True
event_properties = _capture_exception_and_update_properties(
self._ph_client,
outputs,
self._distinct_id,
self._groups,
event_properties,
)
elif outputs is not None:
event_properties["$ai_output_state"] = with_privacy_mode(
self._ph_client, self._privacy_mode, outputs
@@ -566,16 +578,33 @@ class CallbackHandler(BaseCallbackHandler):
"$ai_framework": "langchain",
}
if isinstance(run.insights_properties, dict):
event_properties.update(run.insights_properties)
if run.tools:
event_properties["$ai_tools"] = run.tools
if self._properties:
event_properties.update(self._properties)
if self._distinct_id is None:
event_properties["$process_person_profile"] = False
if isinstance(output, BaseException):
event_properties["$ai_http_status"] = _get_http_status(output)
event_properties["$ai_error"] = _stringify_exception(output)
event_properties["$ai_is_error"] = True
event_properties = _capture_exception_and_update_properties(
self._ph_client,
output,
self._distinct_id,
self._groups,
event_properties,
)
else:
# Add usage
usage = _parse_usage(output)
usage = _parse_usage(output, run.provider, run.model)
event_properties["$ai_input_tokens"] = usage.input_tokens
event_properties["$ai_output_tokens"] = usage.output_tokens
event_properties["$ai_cache_creation_input_tokens"] = (
@@ -600,12 +629,6 @@ class CallbackHandler(BaseCallbackHandler):
self._ph_client, self._privacy_mode, completions
)
if self._properties:
event_properties.update(self._properties)
if self._distinct_id is None:
event_properties["$process_person_profile"] = False
self._ph_client.capture(
distinct_id=self._distinct_id or trace_id,
event="$ai_generation",
@@ -696,6 +719,8 @@ class ModelUsage:
def _parse_usage_model(
usage: Union[BaseModel, dict],
provider: Optional[str] = None,
model: Optional[str] = None,
) -> ModelUsage:
if isinstance(usage, BaseModel):
usage = usage.__dict__
@@ -764,16 +789,32 @@ def _parse_usage_model(
for mapped_key, dataclass_key in field_mapping.items()
},
)
# In LangChain, input_tokens is the sum of input and cache read tokens.
# Our cost calculation expects them to be separate, for Anthropic.
if normalized_usage.input_tokens and normalized_usage.cache_read_tokens:
normalized_usage.input_tokens = max(
normalized_usage.input_tokens - normalized_usage.cache_read_tokens, 0
# For Anthropic providers, LangChain reports input_tokens as the sum of all input tokens.
# Our cost calculation expects them to be separate for Anthropic, so we subtract cache tokens.
# Both cache_read and cache_write tokens should be subtracted since Anthropic's raw API
# reports input_tokens as tokens NOT read from or used to create a cache.
# For other providers (OpenAI, etc.), input_tokens already excludes cache tokens as expected.
# Match logic consistent with plugin-server: exact match on provider OR substring match on model
is_anthropic = False
if provider and provider.lower() == "anthropic":
is_anthropic = True
elif model and "anthropic" in model.lower():
is_anthropic = True
if is_anthropic and normalized_usage.input_tokens:
cache_tokens = (normalized_usage.cache_read_tokens or 0) + (
normalized_usage.cache_write_tokens or 0
)
if cache_tokens > 0:
normalized_usage.input_tokens = max(
normalized_usage.input_tokens - cache_tokens, 0
)
return normalized_usage
def _parse_usage(response: LLMResult) -> ModelUsage:
def _parse_usage(
response: LLMResult, provider: Optional[str] = None, model: Optional[str] = None
) -> ModelUsage:
# langchain-anthropic uses the usage field
llm_usage_keys = ["token_usage", "usage"]
llm_usage: ModelUsage = ModelUsage(
@@ -787,13 +828,15 @@ def _parse_usage(response: LLMResult) -> ModelUsage:
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])
llm_usage = _parse_usage_model(
response.llm_output[key], provider, model
)
break
if hasattr(response, "generations"):
for generation in response.generations:
if "usage" in generation:
llm_usage = _parse_usage_model(generation["usage"])
llm_usage = _parse_usage_model(generation["usage"], provider, model)
break
for generation_chunk in generation:
@@ -801,7 +844,9 @@ def _parse_usage(response: LLMResult) -> ModelUsage:
"usage_metadata" in generation_chunk.generation_info
):
llm_usage = _parse_usage_model(
generation_chunk.generation_info["usage_metadata"]
generation_chunk.generation_info["usage_metadata"],
provider,
model,
)
break
@@ -828,12 +873,33 @@ def _parse_usage(response: LLMResult) -> ModelUsage:
bedrock_anthropic_usage or bedrock_titan_usage or ollama_usage
)
if chunk_usage:
llm_usage = _parse_usage_model(chunk_usage)
llm_usage = _parse_usage_model(chunk_usage, provider, model)
break
return llm_usage
def _capture_exception_and_update_properties(
client: Client,
exception: BaseException,
distinct_id: Optional[Union[str, int, UUID]],
groups: Optional[Dict[str, Any]],
event_properties: Dict[str, Any],
):
if client.enable_exception_autocapture:
exception_id = client.capture_exception(
exception,
distinct_id=distinct_id,
groups=groups,
properties=event_properties,
)
if exception_id:
event_properties["$exception_event_id"] = exception_id
return event_properties
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
@@ -2,7 +2,7 @@ import time
import uuid
from typing import Any, Dict, List, Optional
from posthog.ai.types import TokenUsage
from hanzo_insights.ai.types import TokenUsage
try:
import openai
@@ -11,40 +11,40 @@ except ImportError:
"Please install the OpenAI SDK to use this feature: 'pip install openai'"
)
from posthog.ai.utils import (
from hanzo_insights.ai.utils import (
call_llm_and_track_usage,
extract_available_tool_calls,
merge_usage_stats,
with_privacy_mode,
)
from posthog.ai.openai.openai_converter import (
from hanzo_insights.ai.openai.openai_converter import (
extract_openai_usage_from_chunk,
extract_openai_content_from_chunk,
extract_openai_tool_calls_from_chunk,
accumulate_openai_tool_calls,
)
from posthog.ai.sanitization import sanitize_openai, sanitize_openai_response
from posthog.client import Client as PostHogClient
from posthog import setup
from hanzo_insights.ai.sanitization import sanitize_openai, sanitize_openai_response
from hanzo_insights.client import Client as InsightsClient
from hanzo_insights import setup
class OpenAI(openai.OpenAI):
"""
A wrapper around the OpenAI SDK that automatically sends LLM usage events to PostHog.
A wrapper around the OpenAI SDK that automatically sends LLM usage events to Insights.
"""
_ph_client: PostHogClient
_ph_client: InsightsClient
def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs):
def __init__(self, insights_client: Optional[InsightsClient] = None, **kwargs):
"""
Args:
api_key: OpenAI API key.
posthog_client: If provided, events will be captured via this client instead of the global `posthog`.
insights_client: If provided, events will be captured via this client instead of the global client.
**openai_config: Any additional keyword args to set on openai (e.g. organization="xxx").
"""
super().__init__(**kwargs)
self._ph_client = posthog_client or setup()
self._ph_client = insights_client or setup()
# Store original objects after parent initialization (only if they exist)
self._original_chat = getattr(self, "chat", None)
@@ -67,7 +67,7 @@ class OpenAI(openai.OpenAI):
class WrappedResponses:
"""Wrapper for OpenAI responses that tracks usage in PostHog."""
"""Wrapper for OpenAI responses that tracks usage in Insights."""
def __init__(self, client: OpenAI, original_responses):
self._client = client
@@ -79,34 +79,34 @@ class WrappedResponses:
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,
insights_distinct_id: Optional[str] = None,
insights_trace_id: Optional[str] = None,
insights_properties: Optional[Dict[str, Any]] = None,
insights_privacy_mode: bool = False,
insights_groups: Optional[Dict[str, Any]] = None,
**kwargs: Any,
):
if posthog_trace_id is None:
posthog_trace_id = str(uuid.uuid4())
if insights_trace_id is None:
insights_trace_id = str(uuid.uuid4())
if kwargs.get("stream", False):
return self._create_streaming(
posthog_distinct_id,
posthog_trace_id,
posthog_properties,
posthog_privacy_mode,
posthog_groups,
insights_distinct_id,
insights_trace_id,
insights_properties,
insights_privacy_mode,
insights_groups,
**kwargs,
)
return call_llm_and_track_usage(
posthog_distinct_id,
insights_distinct_id,
self._client._ph_client,
"openai",
posthog_trace_id,
posthog_properties,
posthog_privacy_mode,
posthog_groups,
insights_trace_id,
insights_properties,
insights_privacy_mode,
insights_groups,
self._client.base_url,
self._original.create,
**kwargs,
@@ -114,24 +114,33 @@ class WrappedResponses:
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]],
insights_distinct_id: Optional[str],
insights_trace_id: Optional[str],
insights_properties: Optional[Dict[str, Any]],
insights_privacy_mode: bool,
insights_groups: Optional[Dict[str, Any]],
**kwargs: Any,
):
start_time = time.time()
usage_stats: TokenUsage = TokenUsage()
final_content = []
model_from_response: Optional[str] = None
response = self._original.create(**kwargs)
def generator():
nonlocal usage_stats
nonlocal final_content # noqa: F824
nonlocal model_from_response
try:
for chunk in response:
# Extract model from response object in chunk (for stored prompts)
if hasattr(chunk, "response") and chunk.response:
if model_from_response is None and hasattr(
chunk.response, "model"
):
model_from_response = chunk.response.model
# Extract usage stats from chunk
chunk_usage = extract_openai_usage_from_chunk(chunk, "responses")
@@ -151,58 +160,63 @@ class WrappedResponses:
latency = end_time - start_time
output = final_content
self._capture_streaming_event(
posthog_distinct_id,
posthog_trace_id,
posthog_properties,
posthog_privacy_mode,
posthog_groups,
insights_distinct_id,
insights_trace_id,
insights_properties,
insights_privacy_mode,
insights_groups,
kwargs,
usage_stats,
latency,
output,
None, # Responses API doesn't have tools
model_from_response,
)
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]],
insights_distinct_id: Optional[str],
insights_trace_id: Optional[str],
insights_properties: Optional[Dict[str, Any]],
insights_privacy_mode: bool,
insights_groups: Optional[Dict[str, Any]],
kwargs: Dict[str, Any],
usage_stats: TokenUsage,
latency: float,
output: Any,
available_tool_calls: Optional[List[Dict[str, Any]]] = None,
model_from_response: Optional[str] = None,
):
from posthog.ai.types import StreamingEventData
from posthog.ai.openai.openai_converter import (
from hanzo_insights.ai.types import StreamingEventData
from hanzo_insights.ai.openai.openai_converter import (
format_openai_streaming_input,
format_openai_streaming_output,
)
from posthog.ai.utils import capture_streaming_event
from hanzo_insights.ai.utils import capture_streaming_event
# Prepare standardized event data
formatted_input = format_openai_streaming_input(kwargs, "responses")
sanitized_input = sanitize_openai_response(formatted_input)
# Use model from kwargs, fallback to model from response
model = kwargs.get("model") or model_from_response or "unknown"
event_data = StreamingEventData(
provider="openai",
model=kwargs.get("model", "unknown"),
model=model,
base_url=str(self._client.base_url),
kwargs=kwargs,
formatted_input=sanitized_input,
formatted_output=format_openai_streaming_output(output, "responses"),
usage_stats=usage_stats,
latency=latency,
distinct_id=posthog_distinct_id,
trace_id=posthog_trace_id,
properties=posthog_properties,
privacy_mode=posthog_privacy_mode,
groups=posthog_groups,
distinct_id=insights_distinct_id,
trace_id=insights_trace_id,
properties=insights_properties,
privacy_mode=insights_privacy_mode,
groups=insights_groups,
)
# Use the common capture function
@@ -210,35 +224,35 @@ class WrappedResponses:
def parse(
self,
posthog_distinct_id: Optional[str] = None,
posthog_trace_id: Optional[str] = None,
posthog_properties: Optional[Dict[str, Any]] = None,
posthog_privacy_mode: bool = False,
posthog_groups: Optional[Dict[str, Any]] = None,
insights_distinct_id: Optional[str] = None,
insights_trace_id: Optional[str] = None,
insights_properties: Optional[Dict[str, Any]] = None,
insights_privacy_mode: bool = False,
insights_groups: Optional[Dict[str, Any]] = None,
**kwargs: Any,
):
"""
Parse structured output using OpenAI's 'responses.parse' method, but also track usage in PostHog.
Parse structured output using OpenAI's 'responses.parse' method, but also track usage in Insights.
Args:
posthog_distinct_id: Optional ID to associate with the usage event.
posthog_trace_id: Optional trace UUID for linking events.
posthog_properties: Optional dictionary of extra properties to include in the event.
posthog_privacy_mode: Whether to anonymize the input and output.
posthog_groups: Optional dictionary of groups to associate with the event.
insights_distinct_id: Optional ID to associate with the usage event.
insights_trace_id: Optional trace UUID for linking events.
insights_properties: Optional dictionary of extra properties to include in the event.
insights_privacy_mode: Whether to anonymize the input and output.
insights_groups: Optional dictionary of groups to associate with the event.
**kwargs: Any additional parameters for the OpenAI Responses Parse API.
Returns:
The response from OpenAI's responses.parse call.
"""
return call_llm_and_track_usage(
posthog_distinct_id,
insights_distinct_id,
self._client._ph_client,
"openai",
posthog_trace_id,
posthog_properties,
posthog_privacy_mode,
posthog_groups,
insights_trace_id,
insights_properties,
insights_privacy_mode,
insights_groups,
self._client.base_url,
self._original.parse,
**kwargs,
@@ -246,7 +260,7 @@ class WrappedResponses:
class WrappedChat:
"""Wrapper for OpenAI chat that tracks usage in PostHog."""
"""Wrapper for OpenAI chat that tracks usage in Insights."""
def __init__(self, client: OpenAI, original_chat):
self._client = client
@@ -262,7 +276,7 @@ class WrappedChat:
class WrappedCompletions:
"""Wrapper for OpenAI chat completions that tracks usage in PostHog."""
"""Wrapper for OpenAI chat completions that tracks usage in Insights."""
def __init__(self, client: OpenAI, original_completions):
self._client = client
@@ -274,34 +288,34 @@ class WrappedCompletions:
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,
insights_distinct_id: Optional[str] = None,
insights_trace_id: Optional[str] = None,
insights_properties: Optional[Dict[str, Any]] = None,
insights_privacy_mode: bool = False,
insights_groups: Optional[Dict[str, Any]] = None,
**kwargs: Any,
):
if posthog_trace_id is None:
posthog_trace_id = str(uuid.uuid4())
if insights_trace_id is None:
insights_trace_id = str(uuid.uuid4())
if kwargs.get("stream", False):
return self._create_streaming(
posthog_distinct_id,
posthog_trace_id,
posthog_properties,
posthog_privacy_mode,
posthog_groups,
insights_distinct_id,
insights_trace_id,
insights_properties,
insights_privacy_mode,
insights_groups,
**kwargs,
)
return call_llm_and_track_usage(
posthog_distinct_id,
insights_distinct_id,
self._client._ph_client,
"openai",
posthog_trace_id,
posthog_properties,
posthog_privacy_mode,
posthog_groups,
insights_trace_id,
insights_properties,
insights_privacy_mode,
insights_groups,
self._client.base_url,
self._original.create,
**kwargs,
@@ -309,17 +323,18 @@ class WrappedCompletions:
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]],
insights_distinct_id: Optional[str],
insights_trace_id: Optional[str],
insights_properties: Optional[Dict[str, Any]],
insights_privacy_mode: bool,
insights_groups: Optional[Dict[str, Any]],
**kwargs: Any,
):
start_time = time.time()
usage_stats: TokenUsage = TokenUsage()
accumulated_content = []
accumulated_tool_calls: Dict[int, Dict[str, Any]] = {}
model_from_response: Optional[str] = None
if "stream_options" not in kwargs:
kwargs["stream_options"] = {}
kwargs["stream_options"]["include_usage"] = True
@@ -329,9 +344,14 @@ class WrappedCompletions:
nonlocal usage_stats
nonlocal accumulated_content # noqa: F824
nonlocal accumulated_tool_calls
nonlocal model_from_response
try:
for chunk in response:
# Extract model from chunk (Chat Completions chunks have model field)
if model_from_response is None and hasattr(chunk, "model"):
model_from_response = chunk.model
# Extract usage stats from chunk
chunk_usage = extract_openai_usage_from_chunk(chunk, "chat")
@@ -365,60 +385,65 @@ class WrappedCompletions:
)
self._capture_streaming_event(
posthog_distinct_id,
posthog_trace_id,
posthog_properties,
posthog_privacy_mode,
posthog_groups,
insights_distinct_id,
insights_trace_id,
insights_properties,
insights_privacy_mode,
insights_groups,
kwargs,
usage_stats,
latency,
accumulated_content,
tool_calls_list,
extract_available_tool_calls("openai", kwargs),
model_from_response,
)
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]],
insights_distinct_id: Optional[str],
insights_trace_id: Optional[str],
insights_properties: Optional[Dict[str, Any]],
insights_privacy_mode: bool,
insights_groups: Optional[Dict[str, Any]],
kwargs: Dict[str, Any],
usage_stats: TokenUsage,
latency: float,
output: Any,
tool_calls: Optional[List[Dict[str, Any]]] = None,
available_tool_calls: Optional[List[Dict[str, Any]]] = None,
model_from_response: Optional[str] = None,
):
from posthog.ai.types import StreamingEventData
from posthog.ai.openai.openai_converter import (
from hanzo_insights.ai.types import StreamingEventData
from hanzo_insights.ai.openai.openai_converter import (
format_openai_streaming_input,
format_openai_streaming_output,
)
from posthog.ai.utils import capture_streaming_event
from hanzo_insights.ai.utils import capture_streaming_event
# Prepare standardized event data
formatted_input = format_openai_streaming_input(kwargs, "chat")
sanitized_input = sanitize_openai(formatted_input)
# Use model from kwargs, fallback to model from response
model = kwargs.get("model") or model_from_response or "unknown"
event_data = StreamingEventData(
provider="openai",
model=kwargs.get("model", "unknown"),
model=model,
base_url=str(self._client.base_url),
kwargs=kwargs,
formatted_input=sanitized_input,
formatted_output=format_openai_streaming_output(output, "chat", tool_calls),
usage_stats=usage_stats,
latency=latency,
distinct_id=posthog_distinct_id,
trace_id=posthog_trace_id,
properties=posthog_properties,
privacy_mode=posthog_privacy_mode,
groups=posthog_groups,
distinct_id=insights_distinct_id,
trace_id=insights_trace_id,
properties=insights_properties,
privacy_mode=insights_privacy_mode,
groups=insights_groups,
)
# Use the common capture function
@@ -426,7 +451,7 @@ class WrappedCompletions:
class WrappedEmbeddings:
"""Wrapper for OpenAI embeddings that tracks usage in PostHog."""
"""Wrapper for OpenAI embeddings that tracks usage in Insights."""
def __init__(self, client: OpenAI, original_embeddings):
self._client = client
@@ -438,30 +463,30 @@ class WrappedEmbeddings:
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,
insights_distinct_id: Optional[str] = None,
insights_trace_id: Optional[str] = None,
insights_properties: Optional[Dict[str, Any]] = None,
insights_privacy_mode: bool = False,
insights_groups: Optional[Dict[str, Any]] = None,
**kwargs: Any,
):
"""
Create an embedding using OpenAI's 'embeddings.create' method, but also track usage in PostHog.
Create an embedding using OpenAI's 'embeddings.create' method, but also track usage in Insights.
Args:
posthog_distinct_id: Optional ID to associate with the usage event.
posthog_trace_id: Optional trace UUID for linking events.
posthog_properties: Optional dictionary of extra properties to include in the event.
posthog_privacy_mode: Whether to anonymize the input and output.
posthog_groups: Optional dictionary of groups to associate with the event.
insights_distinct_id: Optional ID to associate with the usage event.
insights_trace_id: Optional trace UUID for linking events.
insights_properties: Optional dictionary of extra properties to include in the event.
insights_privacy_mode: Whether to anonymize the input and output.
insights_groups: Optional dictionary of groups to associate with the event.
**kwargs: Any additional parameters for the OpenAI Embeddings API.
Returns:
The response from OpenAI's embeddings.create call.
"""
if posthog_trace_id is None:
posthog_trace_id = str(uuid.uuid4())
if insights_trace_id is None:
insights_trace_id = str(uuid.uuid4())
start_time = time.time()
response = self._original.create(**kwargs)
@@ -483,34 +508,34 @@ class WrappedEmbeddings:
"$ai_model": kwargs.get("model"),
"$ai_input": with_privacy_mode(
self._client._ph_client,
posthog_privacy_mode,
insights_privacy_mode,
sanitize_openai_response(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_trace_id": insights_trace_id,
"$ai_base_url": str(self._client.base_url),
**(posthog_properties or {}),
**(insights_properties or {}),
}
if posthog_distinct_id is None:
if insights_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,
distinct_id=insights_distinct_id or insights_trace_id,
event="$ai_embedding",
properties=event_properties,
groups=posthog_groups,
groups=insights_groups,
)
return response
class WrappedBeta:
"""Wrapper for OpenAI beta features that tracks usage in PostHog."""
"""Wrapper for OpenAI beta features that tracks usage in Insights."""
def __init__(self, client: OpenAI, original_beta):
self._client = client
@@ -526,7 +551,7 @@ class WrappedBeta:
class WrappedBetaChat:
"""Wrapper for OpenAI beta chat that tracks usage in PostHog."""
"""Wrapper for OpenAI beta chat that tracks usage in Insights."""
def __init__(self, client: OpenAI, original_beta_chat):
self._client = client
@@ -542,7 +567,7 @@ class WrappedBetaChat:
class WrappedBetaCompletions:
"""Wrapper for OpenAI beta chat completions that tracks usage in PostHog."""
"""Wrapper for OpenAI beta chat completions that tracks usage in Insights."""
def __init__(self, client: OpenAI, original_beta_completions):
self._client = client
@@ -554,21 +579,21 @@ class WrappedBetaCompletions:
def parse(
self,
posthog_distinct_id: Optional[str] = None,
posthog_trace_id: Optional[str] = None,
posthog_properties: Optional[Dict[str, Any]] = None,
posthog_privacy_mode: bool = False,
posthog_groups: Optional[Dict[str, Any]] = None,
insights_distinct_id: Optional[str] = None,
insights_trace_id: Optional[str] = None,
insights_properties: Optional[Dict[str, Any]] = None,
insights_privacy_mode: bool = False,
insights_groups: Optional[Dict[str, Any]] = None,
**kwargs: Any,
):
return call_llm_and_track_usage(
posthog_distinct_id,
insights_distinct_id,
self._client._ph_client,
"openai",
posthog_trace_id,
posthog_properties,
posthog_privacy_mode,
posthog_groups,
insights_trace_id,
insights_properties,
insights_privacy_mode,
insights_groups,
self._client.base_url,
self._original.parse,
**kwargs,
@@ -2,7 +2,7 @@ import time
import uuid
from typing import Any, Dict, List, Optional
from posthog.ai.types import TokenUsage
from hanzo_insights.ai.types import TokenUsage
try:
import openai
@@ -11,43 +11,43 @@ except ImportError:
"Please install the OpenAI SDK to use this feature: 'pip install openai'"
)
from posthog import setup
from posthog.ai.utils import (
from hanzo_insights import setup
from hanzo_insights.ai.utils import (
call_llm_and_track_usage_async,
extract_available_tool_calls,
get_model_params,
merge_usage_stats,
with_privacy_mode,
)
from posthog.ai.openai.openai_converter import (
from hanzo_insights.ai.openai.openai_converter import (
extract_openai_usage_from_chunk,
extract_openai_content_from_chunk,
extract_openai_tool_calls_from_chunk,
accumulate_openai_tool_calls,
format_openai_streaming_output,
)
from posthog.ai.sanitization import sanitize_openai, sanitize_openai_response
from posthog.client import Client as PostHogClient
from hanzo_insights.ai.sanitization import sanitize_openai, sanitize_openai_response
from hanzo_insights.client import Client as InsightsClient
class AsyncOpenAI(openai.AsyncOpenAI):
"""
An async wrapper around the OpenAI SDK that automatically sends LLM usage events to PostHog.
An async wrapper around the OpenAI SDK that automatically sends LLM usage events to Insights.
"""
_ph_client: PostHogClient
_ph_client: InsightsClient
def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs):
def __init__(self, insights_client: Optional[InsightsClient] = None, **kwargs):
"""
Args:
api_key: OpenAI API key.
posthog_client: If provided, events will be captured via this client instead
of the global posthog.
insights_client: If provided, events will be captured via this client instead
of the global hanzo_insights.
**openai_config: Any additional keyword args to set on openai (e.g. organization="xxx").
"""
super().__init__(**kwargs)
self._ph_client = posthog_client or setup()
self._ph_client = insights_client or setup()
# Store original objects after parent initialization (only if they exist)
self._original_chat = getattr(self, "chat", None)
@@ -70,7 +70,7 @@ class AsyncOpenAI(openai.AsyncOpenAI):
class WrappedResponses:
"""Async wrapper for OpenAI responses that tracks usage in PostHog."""
"""Async wrapper for OpenAI responses that tracks usage in Insights."""
def __init__(self, client: AsyncOpenAI, original_responses):
self._client = client
@@ -83,34 +83,34 @@ class WrappedResponses:
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,
insights_distinct_id: Optional[str] = None,
insights_trace_id: Optional[str] = None,
insights_properties: Optional[Dict[str, Any]] = None,
insights_privacy_mode: bool = False,
insights_groups: Optional[Dict[str, Any]] = None,
**kwargs: Any,
):
if posthog_trace_id is None:
posthog_trace_id = str(uuid.uuid4())
if insights_trace_id is None:
insights_trace_id = str(uuid.uuid4())
if kwargs.get("stream", False):
return await self._create_streaming(
posthog_distinct_id,
posthog_trace_id,
posthog_properties,
posthog_privacy_mode,
posthog_groups,
insights_distinct_id,
insights_trace_id,
insights_properties,
insights_privacy_mode,
insights_groups,
**kwargs,
)
return await call_llm_and_track_usage_async(
posthog_distinct_id,
insights_distinct_id,
self._client._ph_client,
"openai",
posthog_trace_id,
posthog_properties,
posthog_privacy_mode,
posthog_groups,
insights_trace_id,
insights_properties,
insights_privacy_mode,
insights_groups,
self._client.base_url,
self._original.create,
**kwargs,
@@ -118,24 +118,33 @@ class WrappedResponses:
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]],
insights_distinct_id: Optional[str],
insights_trace_id: Optional[str],
insights_properties: Optional[Dict[str, Any]],
insights_privacy_mode: bool,
insights_groups: Optional[Dict[str, Any]],
**kwargs: Any,
):
start_time = time.time()
usage_stats: TokenUsage = TokenUsage()
final_content = []
model_from_response: Optional[str] = None
response = await self._original.create(**kwargs)
async def async_generator():
nonlocal usage_stats
nonlocal final_content # noqa: F824
nonlocal model_from_response
try:
async for chunk in response:
# Extract model from response object in chunk (for stored prompts)
if hasattr(chunk, "response") and chunk.response:
if model_from_response is None and hasattr(
chunk.response, "model"
):
model_from_response = chunk.response.model
# Extract usage stats from chunk
chunk_usage = extract_openai_usage_from_chunk(chunk, "responses")
@@ -156,48 +165,53 @@ class WrappedResponses:
output = final_content
await self._capture_streaming_event(
posthog_distinct_id,
posthog_trace_id,
posthog_properties,
posthog_privacy_mode,
posthog_groups,
insights_distinct_id,
insights_trace_id,
insights_properties,
insights_privacy_mode,
insights_groups,
kwargs,
usage_stats,
latency,
output,
extract_available_tool_calls("openai", kwargs),
model_from_response,
)
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]],
insights_distinct_id: Optional[str],
insights_trace_id: Optional[str],
insights_properties: Optional[Dict[str, Any]],
insights_privacy_mode: bool,
insights_groups: Optional[Dict[str, Any]],
kwargs: Dict[str, Any],
usage_stats: TokenUsage,
latency: float,
output: Any,
available_tool_calls: Optional[List[Dict[str, Any]]] = None,
model_from_response: Optional[str] = None,
):
if posthog_trace_id is None:
posthog_trace_id = str(uuid.uuid4())
if insights_trace_id is None:
insights_trace_id = str(uuid.uuid4())
# Use model from kwargs, fallback to model from response
model = kwargs.get("model") or model_from_response or "unknown"
event_properties = {
"$ai_provider": "openai",
"$ai_model": kwargs.get("model"),
"$ai_model": model,
"$ai_model_parameters": get_model_params(kwargs),
"$ai_input": with_privacy_mode(
self._client._ph_client,
posthog_privacy_mode,
insights_privacy_mode,
sanitize_openai_response(kwargs.get("input")),
),
"$ai_output_choices": with_privacy_mode(
self._client._ph_client,
posthog_privacy_mode,
insights_privacy_mode,
format_openai_streaming_output(output, "responses"),
),
"$ai_http_status": 200,
@@ -208,9 +222,9 @@ class WrappedResponses:
),
"$ai_reasoning_tokens": usage_stats.get("reasoning_tokens", 0),
"$ai_latency": latency,
"$ai_trace_id": posthog_trace_id,
"$ai_trace_id": insights_trace_id,
"$ai_base_url": str(self._client.base_url),
**(posthog_properties or {}),
**(insights_properties or {}),
}
# Add web search count if present
@@ -225,48 +239,48 @@ class WrappedResponses:
if available_tool_calls:
event_properties["$ai_tools"] = available_tool_calls
if posthog_distinct_id is None:
if insights_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,
distinct_id=insights_distinct_id or insights_trace_id,
event="$ai_generation",
properties=event_properties,
groups=posthog_groups,
groups=insights_groups,
)
async def parse(
self,
posthog_distinct_id: Optional[str] = None,
posthog_trace_id: Optional[str] = None,
posthog_properties: Optional[Dict[str, Any]] = None,
posthog_privacy_mode: bool = False,
posthog_groups: Optional[Dict[str, Any]] = None,
insights_distinct_id: Optional[str] = None,
insights_trace_id: Optional[str] = None,
insights_properties: Optional[Dict[str, Any]] = None,
insights_privacy_mode: bool = False,
insights_groups: Optional[Dict[str, Any]] = None,
**kwargs: Any,
):
"""
Parse structured output using OpenAI's 'responses.parse' method, but also track usage in PostHog.
Parse structured output using OpenAI's 'responses.parse' method, but also track usage in Insights.
Args:
posthog_distinct_id: Optional ID to associate with the usage event.
posthog_trace_id: Optional trace UUID for linking events.
posthog_properties: Optional dictionary of extra properties to include in the event.
posthog_privacy_mode: Whether to anonymize the input and output.
posthog_groups: Optional dictionary of groups to associate with the event.
insights_distinct_id: Optional ID to associate with the usage event.
insights_trace_id: Optional trace UUID for linking events.
insights_properties: Optional dictionary of extra properties to include in the event.
insights_privacy_mode: Whether to anonymize the input and output.
insights_groups: Optional dictionary of groups to associate with the event.
**kwargs: Any additional parameters for the OpenAI Responses Parse API.
Returns:
The response from OpenAI's responses.parse call.
"""
return await call_llm_and_track_usage_async(
posthog_distinct_id,
insights_distinct_id,
self._client._ph_client,
"openai",
posthog_trace_id,
posthog_properties,
posthog_privacy_mode,
posthog_groups,
insights_trace_id,
insights_properties,
insights_privacy_mode,
insights_groups,
self._client.base_url,
self._original.parse,
**kwargs,
@@ -274,7 +288,7 @@ class WrappedResponses:
class WrappedChat:
"""Async wrapper for OpenAI chat that tracks usage in PostHog."""
"""Async wrapper for OpenAI chat that tracks usage in Insights."""
def __init__(self, client: AsyncOpenAI, original_chat):
self._client = client
@@ -290,7 +304,7 @@ class WrappedChat:
class WrappedCompletions:
"""Async wrapper for OpenAI chat completions that tracks usage in PostHog."""
"""Async wrapper for OpenAI chat completions that tracks usage in Insights."""
def __init__(self, client: AsyncOpenAI, original_completions):
self._client = client
@@ -302,35 +316,35 @@ class WrappedCompletions:
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,
insights_distinct_id: Optional[str] = None,
insights_trace_id: Optional[str] = None,
insights_properties: Optional[Dict[str, Any]] = None,
insights_privacy_mode: bool = False,
insights_groups: Optional[Dict[str, Any]] = None,
**kwargs: Any,
):
if posthog_trace_id is None:
posthog_trace_id = str(uuid.uuid4())
if insights_trace_id is None:
insights_trace_id = str(uuid.uuid4())
# If streaming, handle streaming specifically
if kwargs.get("stream", False):
return await self._create_streaming(
posthog_distinct_id,
posthog_trace_id,
posthog_properties,
posthog_privacy_mode,
posthog_groups,
insights_distinct_id,
insights_trace_id,
insights_properties,
insights_privacy_mode,
insights_groups,
**kwargs,
)
response = await call_llm_and_track_usage_async(
posthog_distinct_id,
insights_distinct_id,
self._client._ph_client,
"openai",
posthog_trace_id,
posthog_properties,
posthog_privacy_mode,
posthog_groups,
insights_trace_id,
insights_properties,
insights_privacy_mode,
insights_groups,
self._client.base_url,
self._original.create,
**kwargs,
@@ -339,17 +353,18 @@ class WrappedCompletions:
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]],
insights_distinct_id: Optional[str],
insights_trace_id: Optional[str],
insights_properties: Optional[Dict[str, Any]],
insights_privacy_mode: bool,
insights_groups: Optional[Dict[str, Any]],
**kwargs: Any,
):
start_time = time.time()
usage_stats: TokenUsage = TokenUsage()
accumulated_content = []
accumulated_tool_calls: Dict[int, Dict[str, Any]] = {}
model_from_response: Optional[str] = None
if "stream_options" not in kwargs:
kwargs["stream_options"] = {}
@@ -360,9 +375,14 @@ class WrappedCompletions:
nonlocal usage_stats
nonlocal accumulated_content # noqa: F824
nonlocal accumulated_tool_calls
nonlocal model_from_response
try:
async for chunk in response:
# Extract model from chunk (Chat Completions chunks have model field)
if model_from_response is None and hasattr(chunk, "model"):
model_from_response = chunk.model
# Extract usage stats from chunk
chunk_usage = extract_openai_usage_from_chunk(chunk, "chat")
if chunk_usage:
@@ -394,50 +414,55 @@ class WrappedCompletions:
)
await self._capture_streaming_event(
posthog_distinct_id,
posthog_trace_id,
posthog_properties,
posthog_privacy_mode,
posthog_groups,
insights_distinct_id,
insights_trace_id,
insights_properties,
insights_privacy_mode,
insights_groups,
kwargs,
usage_stats,
latency,
accumulated_content,
tool_calls_list,
extract_available_tool_calls("openai", kwargs),
model_from_response,
)
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]],
insights_distinct_id: Optional[str],
insights_trace_id: Optional[str],
insights_properties: Optional[Dict[str, Any]],
insights_privacy_mode: bool,
insights_groups: Optional[Dict[str, Any]],
kwargs: Dict[str, Any],
usage_stats: TokenUsage,
latency: float,
output: Any,
tool_calls: Optional[List[Dict[str, Any]]] = None,
available_tool_calls: Optional[List[Dict[str, Any]]] = None,
model_from_response: Optional[str] = None,
):
if posthog_trace_id is None:
posthog_trace_id = str(uuid.uuid4())
if insights_trace_id is None:
insights_trace_id = str(uuid.uuid4())
# Use model from kwargs, fallback to model from response
model = kwargs.get("model") or model_from_response or "unknown"
event_properties = {
"$ai_provider": "openai",
"$ai_model": kwargs.get("model"),
"$ai_model": model,
"$ai_model_parameters": get_model_params(kwargs),
"$ai_input": with_privacy_mode(
self._client._ph_client,
posthog_privacy_mode,
insights_privacy_mode,
sanitize_openai(kwargs.get("messages")),
),
"$ai_output_choices": with_privacy_mode(
self._client._ph_client,
posthog_privacy_mode,
insights_privacy_mode,
format_openai_streaming_output(output, "chat", tool_calls),
),
"$ai_http_status": 200,
@@ -448,9 +473,9 @@ class WrappedCompletions:
),
"$ai_reasoning_tokens": usage_stats.get("reasoning_tokens", 0),
"$ai_latency": latency,
"$ai_trace_id": posthog_trace_id,
"$ai_trace_id": insights_trace_id,
"$ai_base_url": str(self._client.base_url),
**(posthog_properties or {}),
**(insights_properties or {}),
}
# Add web search count if present
@@ -466,20 +491,20 @@ class WrappedCompletions:
if available_tool_calls:
event_properties["$ai_tools"] = available_tool_calls
if posthog_distinct_id is None:
if insights_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,
distinct_id=insights_distinct_id or insights_trace_id,
event="$ai_generation",
properties=event_properties,
groups=posthog_groups,
groups=insights_groups,
)
class WrappedEmbeddings:
"""Async wrapper for OpenAI embeddings that tracks usage in PostHog."""
"""Async wrapper for OpenAI embeddings that tracks usage in Insights."""
def __init__(self, client: AsyncOpenAI, original_embeddings):
self._client = client
@@ -492,30 +517,30 @@ class WrappedEmbeddings:
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,
insights_distinct_id: Optional[str] = None,
insights_trace_id: Optional[str] = None,
insights_properties: Optional[Dict[str, Any]] = None,
insights_privacy_mode: bool = False,
insights_groups: Optional[Dict[str, Any]] = None,
**kwargs: Any,
):
"""
Create an embedding using OpenAI's 'embeddings.create' method, but also track usage in PostHog.
Create an embedding using OpenAI's 'embeddings.create' method, but also track usage in Insights.
Args:
posthog_distinct_id: Optional ID to associate with the usage event.
posthog_trace_id: Optional trace UUID for linking events.
posthog_properties: Optional dictionary of extra properties to include in the event.
posthog_privacy_mode: Whether to anonymize the input and output.
posthog_groups: Optional dictionary of groups to associate with the event.
insights_distinct_id: Optional ID to associate with the usage event.
insights_trace_id: Optional trace UUID for linking events.
insights_properties: Optional dictionary of extra properties to include in the event.
insights_privacy_mode: Whether to anonymize the input and output.
insights_groups: Optional dictionary of groups to associate with the event.
**kwargs: Any additional parameters for the OpenAI Embeddings API.
Returns:
The response from OpenAI's embeddings.create call.
"""
if posthog_trace_id is None:
posthog_trace_id = str(uuid.uuid4())
if insights_trace_id is None:
insights_trace_id = str(uuid.uuid4())
start_time = time.time()
response = await self._original.create(**kwargs)
@@ -538,34 +563,34 @@ class WrappedEmbeddings:
"$ai_model": kwargs.get("model"),
"$ai_input": with_privacy_mode(
self._client._ph_client,
posthog_privacy_mode,
insights_privacy_mode,
sanitize_openai_response(kwargs.get("input")),
),
"$ai_http_status": 200,
"$ai_input_tokens": usage_stats.get("input_tokens", 0),
"$ai_latency": latency,
"$ai_trace_id": posthog_trace_id,
"$ai_trace_id": insights_trace_id,
"$ai_base_url": str(self._client.base_url),
**(posthog_properties or {}),
**(insights_properties or {}),
}
if posthog_distinct_id is None:
if insights_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,
distinct_id=insights_distinct_id or insights_trace_id,
event="$ai_embedding",
properties=event_properties,
groups=posthog_groups,
groups=insights_groups,
)
return response
class WrappedBeta:
"""Async wrapper for OpenAI beta features that tracks usage in PostHog."""
"""Async wrapper for OpenAI beta features that tracks usage in Insights."""
def __init__(self, client: AsyncOpenAI, original_beta):
self._client = client
@@ -582,7 +607,7 @@ class WrappedBeta:
class WrappedBetaChat:
"""Async wrapper for OpenAI beta chat that tracks usage in PostHog."""
"""Async wrapper for OpenAI beta chat that tracks usage in Insights."""
def __init__(self, client: AsyncOpenAI, original_beta_chat):
self._client = client
@@ -599,7 +624,7 @@ class WrappedBetaChat:
class WrappedBetaCompletions:
"""Async wrapper for OpenAI beta chat completions that tracks usage in PostHog."""
"""Async wrapper for OpenAI beta chat completions that tracks usage in Insights."""
def __init__(self, client: AsyncOpenAI, original_beta_completions):
self._client = client
@@ -612,21 +637,21 @@ class WrappedBetaCompletions:
async def parse(
self,
posthog_distinct_id: Optional[str] = None,
posthog_trace_id: Optional[str] = None,
posthog_properties: Optional[Dict[str, Any]] = None,
posthog_privacy_mode: bool = False,
posthog_groups: Optional[Dict[str, Any]] = None,
insights_distinct_id: Optional[str] = None,
insights_trace_id: Optional[str] = None,
insights_properties: Optional[Dict[str, Any]] = None,
insights_privacy_mode: bool = False,
insights_groups: Optional[Dict[str, Any]] = None,
**kwargs: Any,
):
return await call_llm_and_track_usage_async(
posthog_distinct_id,
insights_distinct_id,
self._client._ph_client,
"openai",
posthog_trace_id,
posthog_properties,
posthog_privacy_mode,
posthog_groups,
insights_trace_id,
insights_properties,
insights_privacy_mode,
insights_groups,
self._client.base_url,
self._original.parse,
**kwargs,
@@ -2,13 +2,13 @@
OpenAI-specific conversion utilities.
This module handles the conversion of OpenAI API responses and inputs
into standardized formats for PostHog tracking. It supports both
into standardized formats for Insights tracking. It supports both
Chat Completions API and Responses API formats.
"""
from typing import Any, Dict, List, Optional
from posthog.ai.types import (
from hanzo_insights.ai.types import (
FormattedContentItem,
FormattedFunctionCall,
FormattedImageContent,
@@ -16,6 +16,7 @@ from posthog.ai.types import (
FormattedTextContent,
TokenUsage,
)
from hanzo_insights.ai.utils import serialize_raw_usage
def format_openai_response(response: Any) -> List[FormattedMessage]:
@@ -67,6 +68,12 @@ def format_openai_response(response: Any) -> List[FormattedMessage]:
}
)
# Handle audio output (gpt-4o-audio-preview)
if hasattr(choice.message, "audio") and choice.message.audio:
# Convert Pydantic model to dict to capture all fields from OpenAI
audio_dict = choice.message.audio.model_dump()
content.append({"type": "audio", **audio_dict})
if content:
output.append(
{
@@ -423,6 +430,12 @@ def extract_openai_usage_from_response(response: Any) -> TokenUsage:
if web_search_count > 0:
result["web_search_count"] = web_search_count
# Capture raw usage metadata for backend processing
# Serialize to dict here in the converter (not in utils)
serialized = serialize_raw_usage(response.usage)
if serialized:
result["raw_usage"] = serialized
return result
@@ -476,6 +489,12 @@ def extract_openai_usage_from_chunk(
chunk.usage.completion_tokens_details.reasoning_tokens
)
# Capture raw usage metadata for backend processing
# Serialize to dict here in the converter (not in utils)
serialized = serialize_raw_usage(chunk.usage)
if serialized:
usage["raw_usage"] = serialized
elif provider_type == "responses":
# For Responses API, usage is only in chunk.response.usage for completed events
if hasattr(chunk, "type") and chunk.type == "response.completed":
@@ -510,6 +529,12 @@ def extract_openai_usage_from_chunk(
if web_search_count > 0:
usage["web_search_count"] = web_search_count
# Capture raw usage metadata for backend processing
# Serialize to dict here in the converter (not in utils)
serialized = serialize_raw_usage(response_usage)
if serialized:
usage["raw_usage"] = serialized
return usage
@@ -728,8 +753,8 @@ def format_openai_streaming_input(
api_type: Either "chat" or "responses"
Returns:
Formatted input ready for PostHog tracking
Formatted input ready for Insights tracking
"""
from posthog.ai.utils import merge_system_prompt
from hanzo_insights.ai.utils import merge_system_prompt
return merge_system_prompt(kwargs, "openai")
@@ -5,39 +5,39 @@ except ImportError:
"Please install the Open AI SDK to use this feature: 'pip install openai'"
)
from posthog.ai.openai.openai import (
from hanzo_insights.ai.openai.openai import (
WrappedBeta,
WrappedChat,
WrappedEmbeddings,
WrappedResponses,
)
from posthog.ai.openai.openai_async import WrappedBeta as AsyncWrappedBeta
from posthog.ai.openai.openai_async import WrappedChat as AsyncWrappedChat
from posthog.ai.openai.openai_async import WrappedEmbeddings as AsyncWrappedEmbeddings
from posthog.ai.openai.openai_async import WrappedResponses as AsyncWrappedResponses
from hanzo_insights.ai.openai.openai_async import WrappedBeta as AsyncWrappedBeta
from hanzo_insights.ai.openai.openai_async import WrappedChat as AsyncWrappedChat
from hanzo_insights.ai.openai.openai_async import WrappedEmbeddings as AsyncWrappedEmbeddings
from hanzo_insights.ai.openai.openai_async import WrappedResponses as AsyncWrappedResponses
from typing import Optional
from posthog.client import Client as PostHogClient
from posthog import setup
from hanzo_insights.client import Client as InsightsClient
from hanzo_insights import setup
class AzureOpenAI(openai.AzureOpenAI):
"""
A wrapper around the Azure OpenAI SDK that automatically sends LLM usage events to PostHog.
A wrapper around the Azure OpenAI SDK that automatically sends LLM usage events to Insights.
"""
_ph_client: PostHogClient
_ph_client: InsightsClient
def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs):
def __init__(self, insights_client: Optional[InsightsClient] = None, **kwargs):
"""
Args:
api_key: Azure OpenAI API key.
posthog_client: If provided, events will be captured via this client instead
of the global posthog.
insights_client: If provided, events will be captured via this client instead
of the global hanzo_insights.
**openai_config: Any additional keyword args to set on Azure OpenAI (e.g. azure_endpoint="xxx").
"""
super().__init__(**kwargs)
self._ph_client = posthog_client or setup()
self._ph_client = insights_client or setup()
# Store original objects after parent initialization (only if they exist)
self._original_chat = getattr(self, "chat", None)
@@ -61,21 +61,21 @@ class AzureOpenAI(openai.AzureOpenAI):
class AsyncAzureOpenAI(openai.AsyncAzureOpenAI):
"""
An async wrapper around the Azure OpenAI SDK that automatically sends LLM usage events to PostHog.
An async wrapper around the Azure OpenAI SDK that automatically sends LLM usage events to Insights.
"""
_ph_client: PostHogClient
_ph_client: InsightsClient
def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs):
def __init__(self, insights_client: Optional[InsightsClient] = None, **kwargs):
"""
Args:
api_key: Azure OpenAI API key.
posthog_client: If provided, events will be captured via this client instead
of the global posthog.
insights_client: If provided, events will be captured via this client instead
of the global hanzo_insights.
**openai_config: Any additional keyword args to set on Azure OpenAI (e.g. azure_endpoint="xxx").
"""
super().__init__(**kwargs)
self._ph_client = posthog_client or setup()
self._ph_client = insights_client or setup()
# Store original objects after parent initialization (only if they exist)
self._original_chat = getattr(self, "chat", None)
@@ -0,0 +1,76 @@
from __future__ import annotations
from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union
if TYPE_CHECKING:
from agents.tracing import Trace
from hanzo_insights.client import Client
try:
import agents # noqa: F401
except ImportError:
raise ModuleNotFoundError(
"Please install the OpenAI Agents SDK to use this feature: 'pip install openai-agents'"
)
from hanzo_insights.ai.openai_agents.processor import InsightsTracingProcessor
__all__ = ["InsightsTracingProcessor", "instrument"]
def instrument(
client: Optional[Client] = None,
distinct_id: Optional[Union[str, Callable[[Trace], Optional[str]]]] = None,
privacy_mode: bool = False,
groups: Optional[Dict[str, Any]] = None,
properties: Optional[Dict[str, Any]] = None,
) -> InsightsTracingProcessor:
"""
One-liner to instrument OpenAI Agents SDK with Hanzo Insights tracing.
This registers an InsightsTracingProcessor with the OpenAI Agents SDK,
automatically capturing traces, spans, and LLM generations.
Args:
client: Optional Insights client instance. If not provided, uses the default client.
distinct_id: Optional distinct ID to associate with all traces.
Can also be a callable that takes a trace and returns a distinct ID.
privacy_mode: If True, redacts input/output content from events.
groups: Optional Insights groups to associate with events.
properties: Optional additional properties to include with all events.
Returns:
InsightsTracingProcessor: The registered processor instance.
Example:
```python
from hanzo_insights.ai.openai_agents import instrument
# Simple setup
instrument(distinct_id="user@example.com")
# With custom properties
instrument(
distinct_id="user@example.com",
privacy_mode=True,
properties={"environment": "production"}
)
# Now run agents as normal - traces automatically sent to Insights
from agents import Agent, Runner
agent = Agent(name="Assistant", instructions="You are helpful.")
result = Runner.run_sync(agent, "Hello!")
```
"""
from agents.tracing import add_trace_processor
processor = InsightsTracingProcessor(
client=client,
distinct_id=distinct_id,
privacy_mode=privacy_mode,
groups=groups,
properties=properties,
)
add_trace_processor(processor)
return processor
@@ -0,0 +1,863 @@
import json
import logging
import time
from datetime import datetime
from typing import Any, Callable, Dict, Optional, Union
from agents.tracing import Span, Trace
from agents.tracing.processor_interface import TracingProcessor
from agents.tracing.span_data import (
AgentSpanData,
CustomSpanData,
FunctionSpanData,
GenerationSpanData,
GuardrailSpanData,
HandoffSpanData,
MCPListToolsSpanData,
ResponseSpanData,
SpeechGroupSpanData,
SpeechSpanData,
TranscriptionSpanData,
)
from hanzo_insights import setup
from hanzo_insights.client import Client
log = logging.getLogger("hanzo_insights")
def _ensure_serializable(obj: Any) -> Any:
"""Ensure an object is JSON-serializable, converting to str as fallback.
Returns the original object if it's already serializable (dict, list, str,
int, etc.), or str(obj) for non-serializable types so that downstream
json.dumps() calls won't fail.
"""
if obj is None:
return None
try:
json.dumps(obj)
return obj
except (TypeError, ValueError):
return str(obj)
def _parse_iso_timestamp(iso_str: Optional[str]) -> Optional[float]:
"""Parse ISO timestamp to Unix timestamp."""
if not iso_str:
return None
try:
dt = datetime.fromisoformat(iso_str.replace("Z", "+00:00"))
return dt.timestamp()
except (ValueError, AttributeError):
return None
class InsightsTracingProcessor(TracingProcessor):
"""
A tracing processor that sends OpenAI Agents SDK traces to Hanzo Insights.
This processor implements the TracingProcessor interface from the OpenAI Agents SDK
and maps agent traces, spans, and generations to Insights LLM analytics events.
Example:
```python
from agents import Agent, Runner
from agents.tracing import add_trace_processor
from hanzo_insights.ai.openai_agents import InsightsTracingProcessor
# Create and register the processor
processor = InsightsTracingProcessor(
distinct_id="user@example.com",
privacy_mode=False,
)
add_trace_processor(processor)
# Run agents as normal - traces automatically sent to Insights
agent = Agent(name="Assistant", instructions="You are helpful.")
result = Runner.run_sync(agent, "Hello!")
```
"""
def __init__(
self,
client: Optional[Client] = None,
distinct_id: Optional[Union[str, Callable[[Trace], Optional[str]]]] = None,
privacy_mode: bool = False,
groups: Optional[Dict[str, Any]] = None,
properties: Optional[Dict[str, Any]] = None,
):
"""
Initialize the Insights tracing processor.
Args:
client: Optional Insights client instance. If not provided, uses the default client.
distinct_id: Either a string distinct ID or a callable that takes a Trace
and returns a distinct ID. If not provided, uses the trace_id.
privacy_mode: If True, redacts input/output content from events.
groups: Optional Insights groups to associate with all events.
properties: Optional additional properties to include with all events.
"""
self._client = client or setup()
self._distinct_id = distinct_id
self._privacy_mode = privacy_mode
self._groups = groups or {}
self._properties = properties or {}
# Track span start times for latency calculation
self._span_start_times: Dict[str, float] = {}
# Track trace metadata for associating with spans
self._trace_metadata: Dict[str, Dict[str, Any]] = {}
# Max entries to prevent unbounded growth if on_span_end/on_trace_end
# is never called (e.g., due to an exception in the Agents SDK).
self._max_tracked_entries = 10000
def _get_distinct_id(self, trace: Optional[Trace]) -> Optional[str]:
"""Resolve the distinct ID for a trace.
Returns the user-provided distinct ID (string or callable result),
or None if no user-provided ID is available. Callers should treat
None as a signal to use a fallback ID in personless mode.
"""
if callable(self._distinct_id):
if trace:
result = self._distinct_id(trace)
if result:
return str(result)
return None
elif self._distinct_id:
return str(self._distinct_id)
return None
def _with_privacy_mode(self, value: Any) -> Any:
"""Apply privacy mode redaction if enabled."""
if self._privacy_mode or (
hasattr(self._client, "privacy_mode") and self._client.privacy_mode
):
return None
return value
def _evict_stale_entries(self) -> None:
"""Evict oldest entries if dicts exceed max size to prevent unbounded growth."""
if len(self._span_start_times) > self._max_tracked_entries:
# Remove oldest entries by start time
sorted_spans = sorted(self._span_start_times.items(), key=lambda x: x[1])
for span_id, _ in sorted_spans[: len(sorted_spans) // 2]:
del self._span_start_times[span_id]
log.debug(
"Evicted stale span start times (exceeded %d entries)",
self._max_tracked_entries,
)
if len(self._trace_metadata) > self._max_tracked_entries:
# Remove half the entries (oldest inserted via dict ordering in Python 3.7+)
keys = list(self._trace_metadata.keys())
for key in keys[: len(keys) // 2]:
del self._trace_metadata[key]
log.debug(
"Evicted stale trace metadata (exceeded %d entries)",
self._max_tracked_entries,
)
def _get_group_id(self, trace_id: str) -> Optional[str]:
"""Get the group_id for a trace from stored metadata."""
if trace_id in self._trace_metadata:
return self._trace_metadata[trace_id].get("group_id")
return None
def _capture_event(
self,
event: str,
properties: Dict[str, Any],
distinct_id: Optional[str] = None,
) -> None:
"""Capture an event to Insights with error handling.
Args:
distinct_id: The resolved distinct ID. When the user didn't provide
one, callers should pass ``user_distinct_id or fallback_id``
(matching the langchain/openai pattern) and separately set
``$process_person_profile`` in properties.
"""
try:
if not hasattr(self._client, "capture") or not callable(
self._client.capture
):
return
final_properties = {
**properties,
**self._properties,
}
self._client.capture(
distinct_id=distinct_id or "unknown",
event=event,
properties=final_properties,
groups=self._groups,
)
except Exception as e:
log.debug(f"Failed to capture Insights event: {e}")
def on_trace_start(self, trace: Trace) -> None:
"""Called when a new trace begins. Stores metadata for spans; the $ai_trace event is emitted in on_trace_end."""
try:
self._evict_stale_entries()
trace_id = trace.trace_id
trace_name = trace.name
group_id = getattr(trace, "group_id", None)
metadata = getattr(trace, "metadata", None)
distinct_id = self._get_distinct_id(trace)
# Store trace metadata for later (used by spans and on_trace_end)
self._trace_metadata[trace_id] = {
"name": trace_name,
"group_id": group_id,
"metadata": metadata,
"distinct_id": distinct_id,
"start_time": time.time(),
}
except Exception as e:
log.debug(f"Error in on_trace_start: {e}")
def on_trace_end(self, trace: Trace) -> None:
"""Called when a trace completes. Emits the $ai_trace event with full metadata."""
try:
trace_id = trace.trace_id
# Pop stored metadata (also cleans up)
trace_info = self._trace_metadata.pop(trace_id, {})
trace_name = trace_info.get("name") or trace.name
group_id = trace_info.get("group_id") or getattr(trace, "group_id", None)
metadata = trace_info.get("metadata") or getattr(trace, "metadata", None)
distinct_id = trace_info.get("distinct_id") or self._get_distinct_id(trace)
# Calculate trace-level latency
start_time = trace_info.get("start_time")
latency = (time.time() - start_time) if start_time else None
properties = {
"$ai_trace_id": trace_id,
"$ai_trace_name": trace_name,
"$ai_provider": "openai",
"$ai_framework": "openai-agents",
}
if latency is not None:
properties["$ai_latency"] = latency
# Include group_id for linking related traces (e.g., conversation threads)
if group_id:
properties["$ai_group_id"] = group_id
# Include trace metadata if present
if metadata:
properties["$ai_trace_metadata"] = _ensure_serializable(metadata)
if distinct_id is None:
properties["$process_person_profile"] = False
self._capture_event(
event="$ai_trace",
distinct_id=distinct_id or trace_id,
properties=properties,
)
except Exception as e:
log.debug(f"Error in on_trace_end: {e}")
def on_span_start(self, span: Span[Any]) -> None:
"""Called when a new span begins."""
try:
self._evict_stale_entries()
span_id = span.span_id
self._span_start_times[span_id] = time.time()
except Exception as e:
log.debug(f"Error in on_span_start: {e}")
def on_span_end(self, span: Span[Any]) -> None:
"""Called when a span completes."""
try:
span_id = span.span_id
trace_id = span.trace_id
parent_id = span.parent_id
span_data = span.span_data
# Calculate latency
start_time = self._span_start_times.pop(span_id, None)
if start_time:
latency = time.time() - start_time
else:
# Fall back to parsing timestamps
started = _parse_iso_timestamp(span.started_at)
ended = _parse_iso_timestamp(span.ended_at)
latency = (ended - started) if (started and ended) else 0
# Get user-provided distinct ID from trace metadata (resolved at trace start).
# None means no user-provided ID — use trace_id as fallback in personless mode,
# matching the langchain/openai pattern: `distinct_id or trace_id`.
trace_info = self._trace_metadata.get(trace_id, {})
distinct_id = trace_info.get("distinct_id") or self._get_distinct_id(None)
# Get group_id from trace metadata for linking
group_id = self._get_group_id(trace_id)
# Get error info if present
error_info = span.error
error_properties = {}
if error_info:
if isinstance(error_info, dict):
error_message = error_info.get("message", str(error_info))
error_type_raw = error_info.get("type", "")
else:
error_message = str(error_info)
error_type_raw = ""
# Categorize error type for cross-provider filtering/alerting
error_type = "unknown"
if (
"ModelBehaviorError" in error_type_raw
or "ModelBehaviorError" in error_message
):
error_type = "model_behavior_error"
elif "UserError" in error_type_raw or "UserError" in error_message:
error_type = "user_error"
elif (
"InputGuardrailTripwireTriggered" in error_type_raw
or "InputGuardrailTripwireTriggered" in error_message
):
error_type = "input_guardrail_triggered"
elif (
"OutputGuardrailTripwireTriggered" in error_type_raw
or "OutputGuardrailTripwireTriggered" in error_message
):
error_type = "output_guardrail_triggered"
elif (
"MaxTurnsExceeded" in error_type_raw
or "MaxTurnsExceeded" in error_message
):
error_type = "max_turns_exceeded"
error_properties = {
"$ai_is_error": True,
"$ai_error": error_message,
"$ai_error_type": error_type,
}
# Personless mode: no user-provided distinct_id, fallback to trace_id
if distinct_id is None:
error_properties["$process_person_profile"] = False
distinct_id = trace_id
# Dispatch based on span data type
if isinstance(span_data, GenerationSpanData):
self._handle_generation_span(
span_data,
trace_id,
span_id,
parent_id,
latency,
distinct_id,
group_id,
error_properties,
)
elif isinstance(span_data, FunctionSpanData):
self._handle_function_span(
span_data,
trace_id,
span_id,
parent_id,
latency,
distinct_id,
group_id,
error_properties,
)
elif isinstance(span_data, AgentSpanData):
self._handle_agent_span(
span_data,
trace_id,
span_id,
parent_id,
latency,
distinct_id,
group_id,
error_properties,
)
elif isinstance(span_data, HandoffSpanData):
self._handle_handoff_span(
span_data,
trace_id,
span_id,
parent_id,
latency,
distinct_id,
group_id,
error_properties,
)
elif isinstance(span_data, GuardrailSpanData):
self._handle_guardrail_span(
span_data,
trace_id,
span_id,
parent_id,
latency,
distinct_id,
group_id,
error_properties,
)
elif isinstance(span_data, ResponseSpanData):
self._handle_response_span(
span_data,
trace_id,
span_id,
parent_id,
latency,
distinct_id,
group_id,
error_properties,
)
elif isinstance(span_data, CustomSpanData):
self._handle_custom_span(
span_data,
trace_id,
span_id,
parent_id,
latency,
distinct_id,
group_id,
error_properties,
)
elif isinstance(
span_data, (TranscriptionSpanData, SpeechSpanData, SpeechGroupSpanData)
):
self._handle_audio_span(
span_data,
trace_id,
span_id,
parent_id,
latency,
distinct_id,
group_id,
error_properties,
)
elif isinstance(span_data, MCPListToolsSpanData):
self._handle_mcp_span(
span_data,
trace_id,
span_id,
parent_id,
latency,
distinct_id,
group_id,
error_properties,
)
else:
# Unknown span type - capture as generic span
self._handle_generic_span(
span_data,
trace_id,
span_id,
parent_id,
latency,
distinct_id,
group_id,
error_properties,
)
except Exception as e:
log.debug(f"Error in on_span_end: {e}")
def _base_properties(
self,
trace_id: str,
span_id: str,
parent_id: Optional[str],
latency: float,
group_id: Optional[str],
error_properties: Dict[str, Any],
) -> Dict[str, Any]:
"""Build the base properties dict shared by all span handlers."""
properties = {
"$ai_trace_id": trace_id,
"$ai_span_id": span_id,
"$ai_parent_id": parent_id,
"$ai_provider": "openai",
"$ai_framework": "openai-agents",
"$ai_latency": latency,
**error_properties,
}
if group_id:
properties["$ai_group_id"] = group_id
return properties
def _handle_generation_span(
self,
span_data: GenerationSpanData,
trace_id: str,
span_id: str,
parent_id: Optional[str],
latency: float,
distinct_id: str,
group_id: Optional[str],
error_properties: Dict[str, Any],
) -> None:
"""Handle LLM generation spans - maps to $ai_generation event."""
# Extract token usage
usage = span_data.usage or {}
input_tokens = usage.get("input_tokens") or usage.get("prompt_tokens") or 0
output_tokens = (
usage.get("output_tokens") or usage.get("completion_tokens") or 0
)
# Extract model config parameters
model_config = span_data.model_config or {}
model_params = {}
for param in [
"temperature",
"max_tokens",
"top_p",
"frequency_penalty",
"presence_penalty",
]:
if param in model_config:
model_params[param] = model_config[param]
properties = {
**self._base_properties(
trace_id, span_id, parent_id, latency, group_id, error_properties
),
"$ai_model": span_data.model,
"$ai_model_parameters": model_params if model_params else None,
"$ai_input": self._with_privacy_mode(_ensure_serializable(span_data.input)),
"$ai_output_choices": self._with_privacy_mode(
_ensure_serializable(span_data.output)
),
"$ai_input_tokens": input_tokens,
"$ai_output_tokens": output_tokens,
"$ai_total_tokens": (input_tokens or 0) + (output_tokens or 0),
}
# Add optional token fields if present
if usage.get("reasoning_tokens"):
properties["$ai_reasoning_tokens"] = usage["reasoning_tokens"]
if usage.get("cache_read_input_tokens"):
properties["$ai_cache_read_input_tokens"] = usage["cache_read_input_tokens"]
if usage.get("cache_creation_input_tokens"):
properties["$ai_cache_creation_input_tokens"] = usage[
"cache_creation_input_tokens"
]
self._capture_event("$ai_generation", properties, distinct_id)
def _handle_function_span(
self,
span_data: FunctionSpanData,
trace_id: str,
span_id: str,
parent_id: Optional[str],
latency: float,
distinct_id: str,
group_id: Optional[str],
error_properties: Dict[str, Any],
) -> None:
"""Handle function/tool call spans - maps to $ai_span event."""
properties = {
**self._base_properties(
trace_id, span_id, parent_id, latency, group_id, error_properties
),
"$ai_span_name": span_data.name,
"$ai_span_type": "tool",
"$ai_input_state": self._with_privacy_mode(
_ensure_serializable(span_data.input)
),
"$ai_output_state": self._with_privacy_mode(
_ensure_serializable(span_data.output)
),
}
if span_data.mcp_data:
properties["$ai_mcp_data"] = _ensure_serializable(span_data.mcp_data)
self._capture_event("$ai_span", properties, distinct_id)
def _handle_agent_span(
self,
span_data: AgentSpanData,
trace_id: str,
span_id: str,
parent_id: Optional[str],
latency: float,
distinct_id: str,
group_id: Optional[str],
error_properties: Dict[str, Any],
) -> None:
"""Handle agent execution spans - maps to $ai_span event."""
properties = {
**self._base_properties(
trace_id, span_id, parent_id, latency, group_id, error_properties
),
"$ai_span_name": span_data.name,
"$ai_span_type": "agent",
}
if span_data.handoffs:
properties["$ai_agent_handoffs"] = span_data.handoffs
if span_data.tools:
properties["$ai_agent_tools"] = span_data.tools
if span_data.output_type:
properties["$ai_agent_output_type"] = span_data.output_type
self._capture_event("$ai_span", properties, distinct_id)
def _handle_handoff_span(
self,
span_data: HandoffSpanData,
trace_id: str,
span_id: str,
parent_id: Optional[str],
latency: float,
distinct_id: str,
group_id: Optional[str],
error_properties: Dict[str, Any],
) -> None:
"""Handle agent handoff spans - maps to $ai_span event."""
properties = {
**self._base_properties(
trace_id, span_id, parent_id, latency, group_id, error_properties
),
"$ai_span_name": f"{span_data.from_agent} -> {span_data.to_agent}",
"$ai_span_type": "handoff",
"$ai_handoff_from_agent": span_data.from_agent,
"$ai_handoff_to_agent": span_data.to_agent,
}
self._capture_event("$ai_span", properties, distinct_id)
def _handle_guardrail_span(
self,
span_data: GuardrailSpanData,
trace_id: str,
span_id: str,
parent_id: Optional[str],
latency: float,
distinct_id: str,
group_id: Optional[str],
error_properties: Dict[str, Any],
) -> None:
"""Handle guardrail execution spans - maps to $ai_span event."""
properties = {
**self._base_properties(
trace_id, span_id, parent_id, latency, group_id, error_properties
),
"$ai_span_name": span_data.name,
"$ai_span_type": "guardrail",
"$ai_guardrail_triggered": span_data.triggered,
}
self._capture_event("$ai_span", properties, distinct_id)
def _handle_response_span(
self,
span_data: ResponseSpanData,
trace_id: str,
span_id: str,
parent_id: Optional[str],
latency: float,
distinct_id: str,
group_id: Optional[str],
error_properties: Dict[str, Any],
) -> None:
"""Handle OpenAI Response API spans - maps to $ai_generation event."""
response = span_data.response
response_id = response.id if response else None
# Try to extract usage from response
usage = getattr(response, "usage", None) if response else None
input_tokens = 0
output_tokens = 0
if usage:
input_tokens = getattr(usage, "input_tokens", 0) or 0
output_tokens = getattr(usage, "output_tokens", 0) or 0
# Try to extract model from response
model = getattr(response, "model", None) if response else None
properties = {
**self._base_properties(
trace_id, span_id, parent_id, latency, group_id, error_properties
),
"$ai_model": model,
"$ai_response_id": response_id,
"$ai_input": self._with_privacy_mode(_ensure_serializable(span_data.input)),
"$ai_input_tokens": input_tokens,
"$ai_output_tokens": output_tokens,
"$ai_total_tokens": input_tokens + output_tokens,
}
# Extract output content from response
if response:
output_items = getattr(response, "output", None)
if output_items:
properties["$ai_output_choices"] = self._with_privacy_mode(
_ensure_serializable(output_items)
)
self._capture_event("$ai_generation", properties, distinct_id)
def _handle_custom_span(
self,
span_data: CustomSpanData,
trace_id: str,
span_id: str,
parent_id: Optional[str],
latency: float,
distinct_id: str,
group_id: Optional[str],
error_properties: Dict[str, Any],
) -> None:
"""Handle custom user-defined spans - maps to $ai_span event."""
properties = {
**self._base_properties(
trace_id, span_id, parent_id, latency, group_id, error_properties
),
"$ai_span_name": span_data.name,
"$ai_span_type": "custom",
"$ai_custom_data": self._with_privacy_mode(
_ensure_serializable(span_data.data)
),
}
self._capture_event("$ai_span", properties, distinct_id)
def _handle_audio_span(
self,
span_data: Union[TranscriptionSpanData, SpeechSpanData, SpeechGroupSpanData],
trace_id: str,
span_id: str,
parent_id: Optional[str],
latency: float,
distinct_id: str,
group_id: Optional[str],
error_properties: Dict[str, Any],
) -> None:
"""Handle audio-related spans (transcription, speech) - maps to $ai_span event."""
span_type = span_data.type # "transcription", "speech", or "speech_group"
properties = {
**self._base_properties(
trace_id, span_id, parent_id, latency, group_id, error_properties
),
"$ai_span_name": span_type,
"$ai_span_type": span_type,
}
# Add model info if available
if hasattr(span_data, "model") and span_data.model:
properties["$ai_model"] = span_data.model
# Add model config if available (pass-through property)
if hasattr(span_data, "model_config") and span_data.model_config:
properties["model_config"] = _ensure_serializable(span_data.model_config)
# Add time to first audio byte for speech spans (pass-through property)
if hasattr(span_data, "first_content_at") and span_data.first_content_at:
properties["first_content_at"] = span_data.first_content_at
# Add audio format info (pass-through properties)
if hasattr(span_data, "input_format"):
properties["audio_input_format"] = span_data.input_format
if hasattr(span_data, "output_format"):
properties["audio_output_format"] = span_data.output_format
# Add text input for TTS
if (
hasattr(span_data, "input")
and span_data.input
and isinstance(span_data.input, str)
):
properties["$ai_input"] = self._with_privacy_mode(span_data.input)
# Don't include audio data (base64) - just metadata
if hasattr(span_data, "output") and isinstance(span_data.output, str):
# For transcription, output is the text
properties["$ai_output_state"] = self._with_privacy_mode(span_data.output)
self._capture_event("$ai_span", properties, distinct_id)
def _handle_mcp_span(
self,
span_data: MCPListToolsSpanData,
trace_id: str,
span_id: str,
parent_id: Optional[str],
latency: float,
distinct_id: str,
group_id: Optional[str],
error_properties: Dict[str, Any],
) -> None:
"""Handle MCP (Model Context Protocol) spans - maps to $ai_span event."""
properties = {
**self._base_properties(
trace_id, span_id, parent_id, latency, group_id, error_properties
),
"$ai_span_name": f"mcp:{span_data.server}",
"$ai_span_type": "mcp_tools",
"$ai_mcp_server": span_data.server,
"$ai_mcp_tools": span_data.result,
}
self._capture_event("$ai_span", properties, distinct_id)
def _handle_generic_span(
self,
span_data: Any,
trace_id: str,
span_id: str,
parent_id: Optional[str],
latency: float,
distinct_id: str,
group_id: Optional[str],
error_properties: Dict[str, Any],
) -> None:
"""Handle unknown span types - maps to $ai_span event."""
span_type = getattr(span_data, "type", "unknown")
properties = {
**self._base_properties(
trace_id, span_id, parent_id, latency, group_id, error_properties
),
"$ai_span_name": span_type,
"$ai_span_type": span_type,
}
# Try to export span data
if hasattr(span_data, "export"):
try:
exported = span_data.export()
properties["$ai_span_data"] = _ensure_serializable(exported)
except Exception:
pass
self._capture_event("$ai_span", properties, distinct_id)
def shutdown(self) -> None:
"""Clean up resources when the application stops."""
try:
self._span_start_times.clear()
self._trace_metadata.clear()
# Flush the Insights client if possible
if hasattr(self._client, "flush") and callable(self._client.flush):
self._client.flush()
except Exception as e:
log.debug(f"Error in shutdown: {e}")
def force_flush(self) -> None:
"""Force immediate processing of any queued events."""
try:
if hasattr(self._client, "flush") and callable(self._client.flush):
self._client.flush()
except Exception as e:
log.debug(f"Error in force_flush: {e}")
+329
View File
@@ -0,0 +1,329 @@
"""
Prompt management for Hanzo Insights AI SDK.
Fetch and compile LLM prompts from Insights with caching and fallback support.
"""
import logging
import re
import time
import urllib.parse
from typing import Any, Dict, Optional, Union
from hanzo_insights.request import USER_AGENT, _get_session
from hanzo_insights.utils import remove_trailing_slash
log = logging.getLogger("hanzo_insights")
APP_ENDPOINT = "https://us.insights.hanzo.ai"
DEFAULT_CACHE_TTL_SECONDS = 300 # 5 minutes
PromptVariables = Dict[str, Union[str, int, float, bool]]
PromptCacheKey = tuple[str, Optional[int]]
class CachedPrompt:
"""Cached prompt with metadata."""
def __init__(self, prompt: str, fetched_at: float):
self.prompt = prompt
self.fetched_at = fetched_at
def _cache_key(name: str, version: Optional[int]) -> PromptCacheKey:
"""Build a cache key for latest or versioned prompt fetches."""
return (name, version)
def _prompt_reference(name: str, version: Optional[int]) -> str:
"""Format a prompt reference for logs and errors."""
label = f'prompt "{name}"'
if version is not None:
return f"{label} version {version}"
return label
def _is_prompt_api_response(data: Any) -> bool:
"""Check if the response is a valid prompt API response."""
return (
isinstance(data, dict)
and "prompt" in data
and isinstance(data.get("prompt"), str)
)
class Prompts:
"""
Fetch and compile LLM prompts from Insights.
Can be initialized with a Insights client or with direct options.
Examples:
```python
from hanzo_insights import Insights
from hanzo_insights.ai.prompts import Prompts
# With Insights client
client = Insights('phc_xxx', host='https://us.insights.hanzo.ai', personal_api_key='phx_xxx')
prompts = Prompts(client)
# Or with direct options (no Insights client needed)
prompts = Prompts(
personal_api_key='phx_xxx',
project_api_key='phc_xxx',
host='https://us.insights.hanzo.ai',
)
# Fetch with caching and fallback
template = prompts.get('support-system-prompt', fallback='You are a helpful assistant.')
# Fetch a specific published version
prompt_v1 = prompts.get('support-system-prompt', version=1)
# Compile with variables
system_prompt = prompts.compile(template, {
'company': 'Acme Corp',
'tier': 'premium',
})
```
"""
def __init__(
self,
client: Optional[Any] = None,
*,
personal_api_key: Optional[str] = None,
project_api_key: Optional[str] = None,
host: Optional[str] = None,
default_cache_ttl_seconds: Optional[int] = None,
):
"""
Initialize Prompts.
Args:
client: Insights client instance (optional if personal_api_key provided)
personal_api_key: Direct personal API key (optional if client provided)
project_api_key: Direct project API key (optional if client provided)
host: Insights host (defaults to app endpoint)
default_cache_ttl_seconds: Default cache TTL (defaults to 300)
"""
self._default_cache_ttl_seconds = (
default_cache_ttl_seconds or DEFAULT_CACHE_TTL_SECONDS
)
self._cache: Dict[PromptCacheKey, CachedPrompt] = {}
if client is not None:
self._personal_api_key = getattr(client, "personal_api_key", None) or ""
self._project_api_key = getattr(client, "api_key", None) or ""
self._host = remove_trailing_slash(
getattr(client, "raw_host", None) or APP_ENDPOINT
)
else:
self._personal_api_key = personal_api_key or ""
self._project_api_key = project_api_key or ""
self._host = remove_trailing_slash(host or APP_ENDPOINT)
def get(
self,
name: str,
*,
cache_ttl_seconds: Optional[int] = None,
fallback: Optional[str] = None,
version: Optional[int] = None,
) -> str:
"""
Fetch a prompt by name from the Insights API.
Caching behavior:
1. If cache is fresh, return cached value
2. If fetch fails and cache exists (stale), return stale cache with warning
3. If fetch fails and fallback provided, return fallback with warning
4. If fetch fails with no cache/fallback, raise exception
Args:
name: The name of the prompt to fetch
cache_ttl_seconds: Cache TTL in seconds (defaults to instance default)
fallback: Fallback prompt to use if fetch fails and no cache available
version: Specific prompt version to fetch. If None, fetches the latest
version
Returns:
The prompt string
Raises:
Exception: If the prompt cannot be fetched and no fallback is available
"""
ttl = (
cache_ttl_seconds
if cache_ttl_seconds is not None
else self._default_cache_ttl_seconds
)
cache_key = _cache_key(name, version)
# Check cache first
cached = self._cache.get(cache_key)
now = time.time()
if cached is not None:
is_fresh = (now - cached.fetched_at) < ttl
if is_fresh:
return cached.prompt
# Try to fetch from API
try:
prompt = self._fetch_prompt_from_api(name, version)
fetched_at = time.time()
# Update cache
self._cache[cache_key] = CachedPrompt(prompt=prompt, fetched_at=fetched_at)
return prompt
except Exception as error:
prompt_reference = _prompt_reference(name, version)
# Fallback order:
# 1. Return stale cache (with warning)
if cached is not None:
log.warning(
"[Insights Prompts] Failed to fetch %s, using stale cache: %s",
prompt_reference,
error,
)
return cached.prompt
# 2. Return fallback (with warning)
if fallback is not None:
log.warning(
"[Insights Prompts] Failed to fetch %s, using fallback: %s",
prompt_reference,
error,
)
return fallback
# 3. Raise error
raise
def compile(self, prompt: str, variables: PromptVariables) -> str:
"""
Replace {{variableName}} placeholders with values.
Unmatched variables are left unchanged.
Supports variable names with hyphens and dots (e.g., user-id, company.name).
Args:
prompt: The prompt template string
variables: Object containing variable values
Returns:
The compiled prompt string
"""
def replace_variable(match: re.Match) -> str:
variable_name = match.group(1)
if variable_name in variables:
return str(variables[variable_name])
return match.group(0)
return re.sub(r"\{\{([\w.-]+)\}\}", replace_variable, prompt)
def clear_cache(
self, name: Optional[str] = None, *, version: Optional[int] = None
) -> None:
"""
Clear cached prompts.
Args:
name: Specific prompt name to clear. If None, clears all cached prompts.
version: Specific prompt version to clear. Requires name.
"""
if version is not None and name is None:
raise ValueError("'version' requires 'name' to be provided")
if name is None:
self._cache.clear()
return
if version is not None:
self._cache.pop(_cache_key(name, version), None)
return
keys_to_clear = [key for key in self._cache if key[0] == name]
for key in keys_to_clear:
self._cache.pop(key, None)
def _fetch_prompt_from_api(self, name: str, version: Optional[int] = None) -> str:
"""
Fetch prompt from Insights API.
Endpoint:
{host}/api/environments/@current/llm_prompts/name/{encoded_name}/
?token={encoded_project_api_key}[&version={version}]
Auth: Bearer {personal_api_key}
Args:
name: The name of the prompt to fetch
version: Specific prompt version to fetch. If None, fetches the latest
Returns:
The prompt string
Raises:
Exception: If the prompt cannot be fetched
"""
if not self._personal_api_key:
raise Exception(
"[Insights Prompts] personal_api_key is required to fetch prompts. "
"Please provide it when initializing the Prompts instance."
)
if not self._project_api_key:
raise Exception(
"[Insights Prompts] project_api_key is required to fetch prompts. "
"Please provide it when initializing the Prompts instance."
)
encoded_name = urllib.parse.quote(name, safe="")
query_params: Dict[str, Union[str, int]] = {"token": self._project_api_key}
if version is not None:
query_params["version"] = version
encoded_query = urllib.parse.urlencode(query_params)
url = f"{self._host}/api/environments/@current/llm_prompts/name/{encoded_name}/?{encoded_query}"
prompt_reference = _prompt_reference(name, version)
prompt_label = prompt_reference[:1].upper() + prompt_reference[1:]
headers = {
"Authorization": f"Bearer {self._personal_api_key}",
"User-Agent": USER_AGENT,
}
response = _get_session().get(url, headers=headers, timeout=10)
if not response.ok:
if response.status_code == 404:
raise Exception(f"[Insights Prompts] {prompt_label} not found")
if response.status_code == 403:
raise Exception(
f"[Insights Prompts] Access denied for {prompt_reference}. "
"Check that your personal_api_key has the correct permissions and the LLM prompts feature is enabled."
)
raise Exception(
f"[Insights Prompts] Failed to fetch {prompt_label}: HTTP {response.status_code}"
)
try:
data = response.json()
except Exception:
raise Exception(
f"[Insights Prompts] Invalid response format for {prompt_label}"
)
if not _is_prompt_api_response(data):
raise Exception(
f"[Insights Prompts] Invalid response format for {prompt_label}"
)
return data["prompt"]
@@ -1,3 +1,4 @@
import os
import re
from typing import Any
from urllib.parse import urlparse
@@ -5,6 +6,15 @@ from urllib.parse import urlparse
REDACTED_IMAGE_PLACEHOLDER = "[base64 image redacted]"
def _is_multimodal_enabled() -> bool:
"""Check if multimodal capture is enabled via environment variable."""
return os.environ.get("_INTERNAL_LLMA_MULTIMODAL", "").lower() in (
"true",
"1",
"yes",
)
def is_base64_data_url(text: str) -> bool:
return re.match(r"^data:([^;]+);base64,", text) is not None
@@ -27,6 +37,9 @@ def is_raw_base64(text: str) -> bool:
def redact_base64_data_url(value: Any) -> Any:
if _is_multimodal_enabled():
return value
if not isinstance(value, str):
return value
@@ -70,6 +83,12 @@ def sanitize_openai_image(item: Any) -> Any:
if not isinstance(item, dict):
return item
if item.get("type") == "input_image" and isinstance(item.get("image_url"), str):
return {
**item,
"image_url": redact_base64_data_url(item["image_url"]),
}
if (
item.get("type") == "image_url"
and isinstance(item.get("image_url"), dict)
@@ -83,6 +102,11 @@ def sanitize_openai_image(item: Any) -> Any:
},
}
if item.get("type") == "audio" and "data" in item:
if _is_multimodal_enabled():
return item
return {**item, "data": REDACTED_IMAGE_PLACEHOLDER}
return item
@@ -100,6 +124,9 @@ def sanitize_openai_response_image(item: Any) -> Any:
def sanitize_anthropic_image(item: Any) -> Any:
if _is_multimodal_enabled():
return item
if not isinstance(item, dict):
return item
@@ -109,8 +136,6 @@ def sanitize_anthropic_image(item: Any) -> Any:
and item["source"].get("type") == "base64"
and "data" in item["source"]
):
# For Anthropic, if the source type is "base64", we should always redact the data
# The provider is explicitly telling us this is base64 data
return {
**item,
"source": {
@@ -123,6 +148,9 @@ def sanitize_anthropic_image(item: Any) -> Any:
def sanitize_gemini_part(part: Any) -> Any:
if _is_multimodal_enabled():
return part
if not isinstance(part, dict):
return part
@@ -131,8 +159,6 @@ def sanitize_gemini_part(part: Any) -> Any:
and isinstance(part["inline_data"], dict)
and "data" in part["inline_data"]
):
# For Gemini, the inline_data structure indicates base64 data
# We should redact any string data in this context
return {
**part,
"inline_data": {
@@ -185,7 +211,9 @@ def sanitize_langchain_image(item: Any) -> Any:
and isinstance(item.get("source"), dict)
and "data" in item["source"]
):
# Anthropic style - raw base64 in structured format, always redact
if _is_multimodal_enabled():
return item
return {
**item,
"source": {
@@ -1,5 +1,5 @@
"""
Common type definitions for PostHog AI SDK.
Common type definitions for Insights AI SDK.
These types are used for formatting messages and responses across different AI providers
(Anthropic, OpenAI, Gemini, etc.) to ensure consistency in tracking and data structure.
@@ -41,10 +41,10 @@ FormattedContentItem = Union[
class FormattedMessage(TypedDict):
"""
Standardized message format for PostHog tracking.
Standardized message format for Insights tracking.
Used across all providers to ensure consistent message structure
when sending events to PostHog.
when sending events to Insights.
"""
role: str
@@ -64,6 +64,7 @@ class TokenUsage(TypedDict, total=False):
cache_creation_input_tokens: Optional[int]
reasoning_tokens: Optional[int]
web_search_count: Optional[int]
raw_usage: Optional[Any] # Raw provider usage metadata for backend processing
class ProviderResponse(TypedDict, total=False):
+757
View File
@@ -0,0 +1,757 @@
import time
import uuid
from typing import Any, Callable, Dict, List, Optional, cast
from hanzo_insights import get_tags, identify_context, new_context, tag, contexts
from hanzo_insights.ai.sanitization import (
sanitize_anthropic,
sanitize_gemini,
sanitize_langchain,
sanitize_openai,
)
from hanzo_insights.ai.types import FormattedMessage, StreamingEventData, TokenUsage
from hanzo_insights.client import Client as InsightsClient
_TOKEN_PROPERTY_KEYS = frozenset(
{
"$ai_input_tokens",
"$ai_output_tokens",
"$ai_cache_read_input_tokens",
"$ai_cache_creation_input_tokens",
"$ai_total_tokens",
"$ai_reasoning_tokens",
}
)
def _get_tokens_source(
sdk_tags: Dict[str, Any], insights_properties: Optional[Dict[str, Any]]
) -> str:
if insights_properties and any(
key in insights_properties for key in _TOKEN_PROPERTY_KEYS
):
return "passthrough"
return "sdk"
def serialize_raw_usage(raw_usage: Any) -> Optional[Dict[str, Any]]:
"""
Convert raw provider usage objects to JSON-serializable dicts.
Handles Pydantic models (OpenAI/Anthropic) and protobuf-like objects (Gemini)
with a fallback chain to ensure we never pass unserializable objects to Insights.
Args:
raw_usage: Raw usage object from provider SDK
Returns:
Plain dict or None if conversion fails
"""
if raw_usage is None:
return None
# Already a dict
if isinstance(raw_usage, dict):
return raw_usage
# Try Pydantic model_dump() (OpenAI/Anthropic)
if hasattr(raw_usage, "model_dump") and callable(raw_usage.model_dump):
try:
return raw_usage.model_dump()
except Exception:
pass
# Try to_dict() (some protobuf objects)
if hasattr(raw_usage, "to_dict") and callable(raw_usage.to_dict):
try:
return raw_usage.to_dict()
except Exception:
pass
# Try __dict__ / vars() for simple objects
try:
return vars(raw_usage)
except Exception:
pass
# Last resort: convert to string representation
# This ensures we always return something rather than failing
try:
return {"_raw": str(raw_usage)}
except Exception:
return None
def merge_usage_stats(
target: TokenUsage, source: TokenUsage, mode: str = "incremental"
) -> None:
"""
Merge streaming usage statistics into target dict, handling None values.
Supports two modes:
- "incremental": Add source values to target (for APIs that report new tokens)
- "cumulative": Replace target with source values (for APIs that report totals)
Args:
target: Dictionary to update with usage stats
source: TokenUsage that may contain None values
mode: Either "incremental" or "cumulative"
"""
if mode == "incremental":
# Add new values to existing totals
source_input = source.get("input_tokens")
if source_input is not None:
current = target.get("input_tokens") or 0
target["input_tokens"] = current + source_input
source_output = source.get("output_tokens")
if source_output is not None:
current = target.get("output_tokens") or 0
target["output_tokens"] = current + source_output
source_cache_read = source.get("cache_read_input_tokens")
if source_cache_read is not None:
current = target.get("cache_read_input_tokens") or 0
target["cache_read_input_tokens"] = current + source_cache_read
source_cache_creation = source.get("cache_creation_input_tokens")
if source_cache_creation is not None:
current = target.get("cache_creation_input_tokens") or 0
target["cache_creation_input_tokens"] = current + source_cache_creation
source_reasoning = source.get("reasoning_tokens")
if source_reasoning is not None:
current = target.get("reasoning_tokens") or 0
target["reasoning_tokens"] = current + source_reasoning
source_web_search = source.get("web_search_count")
if source_web_search is not None:
current = target.get("web_search_count") or 0
target["web_search_count"] = max(current, source_web_search)
# Merge raw_usage to avoid losing data from earlier events
# For Anthropic streaming: message_start has input tokens, message_delta has output
# Note: raw_usage is already serialized by converters, so it's a dict
source_raw_usage = source.get("raw_usage")
if source_raw_usage is not None and isinstance(source_raw_usage, dict):
current_raw_value = target.get("raw_usage")
current_raw: Dict[str, Any] = (
current_raw_value if isinstance(current_raw_value, dict) else {}
)
target["raw_usage"] = {**current_raw, **source_raw_usage}
elif mode == "cumulative":
# Replace with latest values (already cumulative)
if source.get("input_tokens") is not None:
target["input_tokens"] = source["input_tokens"]
if source.get("output_tokens") is not None:
target["output_tokens"] = source["output_tokens"]
if source.get("cache_read_input_tokens") is not None:
target["cache_read_input_tokens"] = source["cache_read_input_tokens"]
if source.get("cache_creation_input_tokens") is not None:
target["cache_creation_input_tokens"] = source[
"cache_creation_input_tokens"
]
if source.get("reasoning_tokens") is not None:
target["reasoning_tokens"] = source["reasoning_tokens"]
if source.get("web_search_count") is not None:
target["web_search_count"] = source["web_search_count"]
# Note: raw_usage is already serialized by converters, so it's a dict
if source.get("raw_usage") is not None:
target["raw_usage"] = source["raw_usage"]
else:
raise ValueError(f"Invalid mode: {mode}. Must be 'incremental' or 'cumulative'")
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) -> TokenUsage:
"""
Extract usage statistics from response based on provider.
Delegates to provider-specific converter functions.
"""
if provider == "anthropic":
from hanzo_insights.ai.anthropic.anthropic_converter import (
extract_anthropic_usage_from_response,
)
return extract_anthropic_usage_from_response(response)
elif provider == "openai":
from hanzo_insights.ai.openai.openai_converter import (
extract_openai_usage_from_response,
)
return extract_openai_usage_from_response(response)
elif provider == "gemini":
from hanzo_insights.ai.gemini.gemini_converter import (
extract_gemini_usage_from_response,
)
return extract_gemini_usage_from_response(response)
return TokenUsage(input_tokens=0, output_tokens=0)
def format_response(response, provider: str):
"""
Format a regular (non-streaming) response.
"""
if provider == "anthropic":
from hanzo_insights.ai.anthropic.anthropic_converter import format_anthropic_response
return format_anthropic_response(response)
elif provider == "openai":
from hanzo_insights.ai.openai.openai_converter import format_openai_response
return format_openai_response(response)
elif provider == "gemini":
from hanzo_insights.ai.gemini.gemini_converter import format_gemini_response
return format_gemini_response(response)
return []
def extract_available_tool_calls(provider: str, kwargs: Dict[str, Any]):
"""
Extract available tool calls for the given provider.
"""
if provider == "anthropic":
from hanzo_insights.ai.anthropic.anthropic_converter import extract_anthropic_tools
return extract_anthropic_tools(kwargs)
elif provider == "gemini":
from hanzo_insights.ai.gemini.gemini_converter import extract_gemini_tools
return extract_gemini_tools(kwargs)
elif provider == "openai":
from hanzo_insights.ai.openai.openai_converter import extract_openai_tools
return extract_openai_tools(kwargs)
return None
def merge_system_prompt(
kwargs: Dict[str, Any], provider: str
) -> List[FormattedMessage]:
"""
Merge system prompts and format messages for the given provider.
"""
if provider == "anthropic":
from hanzo_insights.ai.anthropic.anthropic_converter import format_anthropic_input
messages = kwargs.get("messages") or []
system = kwargs.get("system")
return format_anthropic_input(messages, system)
elif provider == "gemini":
from hanzo_insights.ai.gemini.gemini_converter import format_gemini_input_with_system
contents = kwargs.get("contents", [])
config = kwargs.get("config")
return format_gemini_input_with_system(contents, config)
elif provider == "openai":
from hanzo_insights.ai.openai.openai_converter import format_openai_input
# For OpenAI, handle both Chat Completions and Responses API
messages_param = kwargs.get("messages")
input_param = kwargs.get("input")
# Get base formatted messages
messages = format_openai_input(messages_param, input_param)
# Check if system prompt is provided as a separate parameter
if kwargs.get("system") is not None:
has_system = any(msg.get("role") == "system" for msg in messages)
if not has_system:
system_msg = cast(
FormattedMessage,
{"role": "system", "content": kwargs.get("system")},
)
messages = [system_msg] + messages
# For Responses API, add instructions to the system prompt if provided
if kwargs.get("instructions") is not None:
# Find the system message if it exists
system_idx = next(
(i for i, msg in enumerate(messages) if msg.get("role") == "system"),
None,
)
if system_idx is not None:
# Append instructions to existing system message
system_content = messages[system_idx].get("content", "")
messages[system_idx]["content"] = (
f"{system_content}\n\n{kwargs.get('instructions')}"
)
else:
# Create a new system message with instructions
instruction_msg = cast(
FormattedMessage,
{"role": "system", "content": kwargs.get("instructions")},
)
messages = [instruction_msg] + messages
return messages
# Default case - return empty list
return []
def call_llm_and_track_usage(
insights_distinct_id: Optional[str],
ph_client: InsightsClient,
provider: str,
insights_trace_id: Optional[str],
insights_properties: Optional[Dict[str, Any]],
insights_privacy_mode: bool,
insights_groups: Optional[Dict[str, Any]],
base_url: str,
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: TokenUsage = TokenUsage()
error_params: Dict[str, Any] = {}
with new_context(client=ph_client, capture_exceptions=False):
if insights_distinct_id:
identify_context(insights_distinct_id)
try:
response = call_method(**kwargs)
except Exception as exc:
error = exc
http_status = getattr(
exc, "status_code", 0
) # default to 0 becuase its likely an SDK error
error_params = {
"$ai_is_error": True,
"$ai_error": exc.__str__(),
}
# TODO: Add exception capture for OpenAI/Anthropic/Gemini wrappers when
# enable_exception_autocapture is True, similar to LangChain callbacks.
# See _capture_exception_and_update_properties in langchain/callbacks.py
finally:
end_time = time.time()
latency = end_time - start_time
if insights_trace_id is None:
insights_trace_id = str(uuid.uuid4())
# Check if we have a real user distinct_id (from param or outer context)
has_person_distinct_id = (
insights_distinct_id is not None
or contexts.get_context_distinct_id() is not None
)
if not has_person_distinct_id:
# Fall back to trace_id as distinct_id when no real user id is available.
identify_context(insights_trace_id)
if response and (
hasattr(response, "usage")
or (provider == "gemini" and hasattr(response, "usage_metadata"))
):
usage = get_usage(response, provider)
messages = merge_system_prompt(kwargs, provider)
sanitized_messages = sanitize_messages(messages, provider)
tag("$ai_provider", provider)
tag("$ai_model", kwargs.get("model") or getattr(response, "model", None))
tag("$ai_model_parameters", get_model_params(kwargs))
tag(
"$ai_input",
with_privacy_mode(ph_client, insights_privacy_mode, sanitized_messages),
)
tag(
"$ai_output_choices",
with_privacy_mode(
ph_client, insights_privacy_mode, format_response(response, provider)
),
)
tag("$ai_http_status", http_status)
tag("$ai_input_tokens", usage.get("input_tokens", 0))
tag("$ai_output_tokens", usage.get("output_tokens", 0))
tag("$ai_latency", latency)
tag("$ai_trace_id", insights_trace_id)
tag("$ai_base_url", str(base_url))
available_tool_calls = extract_available_tool_calls(provider, kwargs)
if available_tool_calls:
tag("$ai_tools", available_tool_calls)
cache_read = usage.get("cache_read_input_tokens")
if cache_read is not None and cache_read > 0:
tag("$ai_cache_read_input_tokens", cache_read)
cache_creation = usage.get("cache_creation_input_tokens")
if cache_creation is not None and cache_creation > 0:
tag("$ai_cache_creation_input_tokens", cache_creation)
reasoning = usage.get("reasoning_tokens")
if reasoning is not None and reasoning > 0:
tag("$ai_reasoning_tokens", reasoning)
web_search_count = usage.get("web_search_count")
if web_search_count is not None and web_search_count > 0:
tag("$ai_web_search_count", web_search_count)
raw_usage = usage.get("raw_usage")
if raw_usage is not None:
# Already serialized by converters
tag("$ai_usage", raw_usage)
if not has_person_distinct_id:
tag("$process_person_profile", False)
# Process instructions for Responses API
if provider == "openai" and kwargs.get("instructions") is not None:
tag(
"$ai_instructions",
with_privacy_mode(
ph_client, insights_privacy_mode, kwargs.get("instructions")
),
)
# send the event to Insights
if hasattr(ph_client, "capture") and callable(ph_client.capture):
sdk_tags = get_tags()
merged_properties = {
**sdk_tags,
**(insights_properties or {}),
**(error_params or {}),
}
merged_properties["$ai_tokens_source"] = _get_tokens_source(
sdk_tags, insights_properties
)
ph_client.capture(
distinct_id=contexts.get_context_distinct_id(),
event="$ai_generation",
properties=merged_properties,
groups=insights_groups,
)
if error:
raise error
return response
async def call_llm_and_track_usage_async(
insights_distinct_id: Optional[str],
ph_client: InsightsClient,
provider: str,
insights_trace_id: Optional[str],
insights_properties: Optional[Dict[str, Any]],
insights_privacy_mode: bool,
insights_groups: Optional[Dict[str, Any]],
base_url: str,
call_async_method: Callable[..., Any],
**kwargs: Any,
) -> Any:
start_time = time.time()
response = None
error = None
http_status = 200
usage: TokenUsage = TokenUsage()
error_params: Dict[str, Any] = {}
with new_context(client=ph_client, capture_exceptions=False):
if insights_distinct_id:
identify_context(insights_distinct_id)
try:
response = await call_async_method(**kwargs)
except Exception as exc:
error = exc
http_status = getattr(
exc, "status_code", 0
) # default to 0 because its likely an SDK error
error_params = {
"$ai_is_error": True,
"$ai_error": exc.__str__(),
}
# TODO: Add exception capture for OpenAI/Anthropic/Gemini wrappers when
# enable_exception_autocapture is True, similar to LangChain callbacks.
# See _capture_exception_and_update_properties in langchain/callbacks.py
finally:
end_time = time.time()
latency = end_time - start_time
if insights_trace_id is None:
insights_trace_id = str(uuid.uuid4())
# Check if we have a real user distinct_id (from param or outer context)
has_person_distinct_id = (
insights_distinct_id is not None
or contexts.get_context_distinct_id() is not None
)
if not has_person_distinct_id:
# Fall back to trace_id as distinct_id when no real user id is available.
identify_context(insights_trace_id)
if response and (
hasattr(response, "usage")
or (provider == "gemini" and hasattr(response, "usage_metadata"))
):
usage = get_usage(response, provider)
messages = merge_system_prompt(kwargs, provider)
sanitized_messages = sanitize_messages(messages, provider)
tag("$ai_provider", provider)
tag("$ai_model", kwargs.get("model") or getattr(response, "model", None))
tag("$ai_model_parameters", get_model_params(kwargs))
tag(
"$ai_input",
with_privacy_mode(ph_client, insights_privacy_mode, sanitized_messages),
)
tag(
"$ai_output_choices",
with_privacy_mode(
ph_client, insights_privacy_mode, format_response(response, provider)
),
)
tag("$ai_http_status", http_status)
tag("$ai_input_tokens", usage.get("input_tokens", 0))
tag("$ai_output_tokens", usage.get("output_tokens", 0))
tag("$ai_latency", latency)
tag("$ai_trace_id", insights_trace_id)
tag("$ai_base_url", str(base_url))
available_tool_calls = extract_available_tool_calls(provider, kwargs)
if available_tool_calls:
tag("$ai_tools", available_tool_calls)
cache_read = usage.get("cache_read_input_tokens")
if cache_read is not None and cache_read > 0:
tag("$ai_cache_read_input_tokens", cache_read)
cache_creation = usage.get("cache_creation_input_tokens")
if cache_creation is not None and cache_creation > 0:
tag("$ai_cache_creation_input_tokens", cache_creation)
reasoning = usage.get("reasoning_tokens")
if reasoning is not None and reasoning > 0:
tag("$ai_reasoning_tokens", reasoning)
web_search_count = usage.get("web_search_count")
if web_search_count is not None and web_search_count > 0:
tag("$ai_web_search_count", web_search_count)
raw_usage = usage.get("raw_usage")
if raw_usage is not None:
# Already serialized by converters
tag("$ai_usage", raw_usage)
if not has_person_distinct_id:
tag("$process_person_profile", False)
# Process instructions for Responses API
if provider == "openai" and kwargs.get("instructions") is not None:
tag(
"$ai_instructions",
with_privacy_mode(
ph_client, insights_privacy_mode, kwargs.get("instructions")
),
)
# send the event to Insights
if hasattr(ph_client, "capture") and callable(ph_client.capture):
sdk_tags = get_tags()
merged_properties = {
**sdk_tags,
**(insights_properties or {}),
**(error_params or {}),
}
merged_properties["$ai_tokens_source"] = _get_tokens_source(
sdk_tags, insights_properties
)
ph_client.capture(
distinct_id=contexts.get_context_distinct_id(),
event="$ai_generation",
properties=merged_properties,
groups=insights_groups,
)
if error:
raise error
return response
def sanitize_messages(data: Any, provider: str) -> Any:
"""Sanitize messages using provider-specific sanitization functions."""
if provider == "anthropic":
return sanitize_anthropic(data)
elif provider == "openai":
return sanitize_openai(data)
elif provider == "gemini":
return sanitize_gemini(data)
elif provider == "langchain":
return sanitize_langchain(data)
return data
def with_privacy_mode(ph_client: InsightsClient, privacy_mode: bool, value: Any):
if ph_client.privacy_mode or privacy_mode:
return None
return value
def capture_streaming_event(
ph_client: InsightsClient,
event_data: StreamingEventData,
):
"""
Unified streaming event capture for all LLM providers.
This function handles the common logic for capturing streaming events across all providers.
All provider-specific formatting should be done BEFORE calling this function.
The function handles:
- Building Insights event properties
- Extracting and adding tools based on provider
- Applying privacy mode
- Adding special token fields (cache, reasoning)
- Provider-specific fields (e.g., OpenAI instructions)
- Sending the event to Insights
Args:
ph_client: Insights client instance
event_data: Standardized streaming event data containing all necessary information
"""
trace_id = event_data.get("trace_id") or str(uuid.uuid4())
# Build base event properties
event_properties = {
"$ai_provider": event_data["provider"],
"$ai_model": event_data["model"],
"$ai_model_parameters": get_model_params(event_data["kwargs"]),
"$ai_input": with_privacy_mode(
ph_client,
event_data["privacy_mode"],
event_data["formatted_input"],
),
"$ai_output_choices": with_privacy_mode(
ph_client,
event_data["privacy_mode"],
event_data["formatted_output"],
),
"$ai_http_status": 200,
"$ai_input_tokens": event_data["usage_stats"].get("input_tokens", 0),
"$ai_output_tokens": event_data["usage_stats"].get("output_tokens", 0),
"$ai_latency": event_data["latency"],
"$ai_trace_id": trace_id,
"$ai_base_url": str(event_data["base_url"]),
**(event_data.get("properties") or {}),
}
# Determine token source: SDK-computed vs externally overridden
sdk_token_tags = {
"$ai_input_tokens": event_data["usage_stats"].get("input_tokens", 0),
"$ai_output_tokens": event_data["usage_stats"].get("output_tokens", 0),
}
event_properties["$ai_tokens_source"] = _get_tokens_source(
sdk_token_tags, event_data.get("properties")
)
# Extract and add tools based on provider
available_tools = extract_available_tool_calls(
event_data["provider"],
event_data["kwargs"],
)
if available_tools:
event_properties["$ai_tools"] = available_tools
# Add optional token fields
# For Anthropic, always include cache fields even if 0 (backward compatibility)
# For others, only include if present and non-zero
if event_data["provider"] == "anthropic":
# Anthropic always includes cache fields
cache_read = event_data["usage_stats"].get("cache_read_input_tokens", 0)
cache_creation = event_data["usage_stats"].get("cache_creation_input_tokens", 0)
event_properties["$ai_cache_read_input_tokens"] = cache_read
event_properties["$ai_cache_creation_input_tokens"] = cache_creation
else:
# Other providers only include if non-zero
optional_token_fields = [
"cache_read_input_tokens",
"cache_creation_input_tokens",
"reasoning_tokens",
]
for field in optional_token_fields:
value = event_data["usage_stats"].get(field)
if value is not None and isinstance(value, int) and value > 0:
event_properties[f"$ai_{field}"] = value
# Add web search count if present (all providers)
web_search_count = event_data["usage_stats"].get("web_search_count")
if (
web_search_count is not None
and isinstance(web_search_count, int)
and web_search_count > 0
):
event_properties["$ai_web_search_count"] = web_search_count
# Add raw usage metadata if present (all providers)
raw_usage = event_data["usage_stats"].get("raw_usage")
if raw_usage is not None:
# Already serialized by converters
event_properties["$ai_usage"] = raw_usage
# Handle provider-specific fields
if (
event_data["provider"] == "openai"
and event_data["kwargs"].get("instructions") is not None
):
event_properties["$ai_instructions"] = with_privacy_mode(
ph_client,
event_data["privacy_mode"],
event_data["kwargs"]["instructions"],
)
if event_data.get("distinct_id") is None:
event_properties["$process_person_profile"] = False
# Send event to Insights
if hasattr(ph_client, "capture"):
ph_client.capture(
distinct_id=event_data.get("distinct_id") or trace_id,
event="$ai_generation",
properties=event_properties,
groups=event_data.get("groups"),
)
+1 -1
View File
@@ -5,7 +5,7 @@ from datetime import datetime
import numbers
from uuid import UUID
from posthog.types import SendFeatureFlagsOptions
from hanzo_insights.types import SendFeatureFlagsOptions
ID_TYPES = Union[numbers.Number, str, UUID, int]
File diff suppressed because it is too large Load Diff
@@ -3,9 +3,7 @@ import logging
import time
from threading import Thread
import backoff
from posthog.request import APIError, DatetimeSerializer, batch_post
from hanzo_insights.request import APIError, DatetimeSerializer, batch_post
try:
from queue import Empty
@@ -23,7 +21,7 @@ BATCH_SIZE_LIMIT = 5 * 1024 * 1024
class Consumer(Thread):
"""Consumes the messages from the client's queue."""
log = logging.getLogger("posthog")
log = logging.getLogger("hanzo_insights")
def __init__(
self,
@@ -84,12 +82,16 @@ class Consumer(Thread):
self.log.error("error uploading: %s", e)
success = False
if self.on_error:
self.on_error(e, batch)
try:
self.on_error(e, batch)
except Exception as e:
self.log.error("on_error handler failed: %s", e)
finally:
# mark items as acknowledged from queue
for item in batch:
self.queue.task_done()
return success
return success
def next(self):
"""Return the next batch of items to upload."""
@@ -124,29 +126,41 @@ class Consumer(Thread):
def request(self, batch):
"""Attempt to upload the batch and retry before raising an error"""
def fatal_exception(exc):
def is_retryable(exc):
if isinstance(exc, APIError):
# retry on server errors and client errors
# with 429 status code (rate limited),
# with 408 (request timeout) or 429 (rate limited),
# don't retry on other client errors
if exc.status == "N/A":
return False
return (400 <= exc.status < 500) and exc.status != 429
return not ((400 <= exc.status < 500) and exc.status not in (408, 429))
else:
# retry on all other errors (eg. network)
return False
return True
@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,
historical_migration=self.historical_migration,
)
last_exc = None
for attempt in range(self.retries + 1):
try:
batch_post(
self.api_key,
self.host,
gzip=self.gzip,
timeout=self.timeout,
batch=batch,
historical_migration=self.historical_migration,
)
return
except Exception as e:
last_exc = e
if not is_retryable(e):
raise
if attempt < self.retries:
# Respect Retry-After header if present, otherwise use exponential backoff
retry_after = getattr(e, "retry_after", None)
if retry_after and retry_after > 0:
time.sleep(retry_after)
else:
time.sleep(min(2**attempt, 30))
send_request()
if last_exc:
raise last_exc
@@ -4,7 +4,7 @@ from typing import Optional, Any, Callable, Dict, TypeVar, cast, TYPE_CHECKING
if TYPE_CHECKING:
# To avoid circular imports
from posthog.client import Client
from hanzo_insights.client import Client
class ContextScope:
@@ -21,6 +21,7 @@ class ContextScope:
self.capture_exceptions = capture_exceptions
self.session_id: Optional[str] = None
self.distinct_id: Optional[str] = None
self.device_id: Optional[str] = None
self.tags: Dict[str, Any] = {}
self.capture_exception_code_variables: Optional[bool] = None
self.code_variables_mask_patterns: Optional[list] = None
@@ -32,6 +33,9 @@ class ContextScope:
def set_distinct_id(self, distinct_id: str):
self.distinct_id = distinct_id
def set_device_id(self, device_id: str):
self.device_id = device_id
def add_tag(self, key: str, value: Any):
self.tags[key] = value
@@ -61,15 +65,21 @@ class ContextScope:
return self.parent.get_distinct_id()
return None
def get_device_id(self) -> Optional[str]:
if self.device_id is not None:
return self.device_id
if self.parent is not None and not self.fresh:
return self.parent.get_device_id()
return None
def collect_tags(self) -> Dict[str, Any]:
tags = self.tags.copy()
if self.parent and not self.fresh:
# We want child tags to take precedence over parent tags,
# so we can't use a simple update here, instead collecting
# the parent tags and then updating with the child tags.
new_tags = self.parent.collect_tags()
tags.update(new_tags)
return tags
# so collect parent tags first, then update with child tags.
tags = self.parent.collect_tags()
tags.update(self.tags)
return tags
return self.tags.copy()
def get_capture_exception_code_variables(self) -> Optional[bool]:
if self.capture_exception_code_variables is not None:
@@ -94,7 +104,7 @@ class ContextScope:
_context_stack: contextvars.ContextVar[Optional[ContextScope]] = contextvars.ContextVar(
"posthog_context_stack", default=None
"insights_context_stack", default=None
)
@@ -124,32 +134,32 @@ def new_context(
If provided, the client will be used to capture exceptions within the context.
If not provided, the default (global) client will be used. Note that the passed
client is only used to capture exceptions within the context - other events captured
within the context via `Client.capture` or `posthog.capture` will still carry the context
within the context via `Client.capture` or `hanzo_insights.capture` will still carry the context
state (tags, identity, session id), but will be captured by the client directly used (or
the global one, in the case of `posthog.capture`)
the global one, in the case of `hanzo_insights.capture`)
Examples:
```python
# Inherit parent context tags
with posthog.new_context():
posthog.tag("request_id", "123")
with hanzo_insights.new_context():
hanzo_insights.tag("request_id", "123")
# Both this event and the exception will be tagged with the context tags
posthog.capture("event_name", {"property": "value"})
hanzo_insights.capture("event_name", {"property": "value"})
raise ValueError("Something went wrong")
```
```python
# Start with fresh context (no inherited tags)
with posthog.new_context(fresh=True):
posthog.tag("request_id", "123")
with hanzo_insights.new_context(fresh=True):
hanzo_insights.tag("request_id", "123")
# Both this event and the exception will be tagged with the context tags
posthog.capture("event_name", {"property": "value"})
hanzo_insights.capture("event_name", {"property": "value"})
raise ValueError("Something went wrong")
```
Category:
Contexts
"""
from posthog import capture_exception
from hanzo_insights import capture_exception
current_context = _get_current_context()
new_context = ContextScope(current_context, fresh, capture_exceptions, client)
@@ -179,7 +189,7 @@ def tag(key: str, value: Any) -> None:
Example:
```python
posthog.tag("user_id", "123")
hanzo_insights.tag("user_id", "123")
```
Category:
@@ -211,7 +221,7 @@ def identify_context(distinct_id: str) -> None:
"""
Identify the current context with a distinct ID, associating all events captured in this or
child contexts with the given distinct ID (unless identify_context is called again). This is overridden by
distinct id's passed directly to posthog.capture and related methods (identify, set etc). Entering a
distinct id's passed directly to hanzo_insights.capture and related methods (identify, set etc). Entering a
fresh context will clear the context-level distinct ID. The distinct-id passed should be uniquely associated
with one of your users. Events captured outside of a context, or in a context with no associated distinct
ID, will be assigned a random UUID, and captured as "personless".
@@ -234,7 +244,7 @@ def set_context_session(session_id: str) -> None:
Entering a fresh context will clear the context-level session ID.
Args:
session_id: The session ID to associate with the current context and its children. See https://posthog.com/docs/data/sessions
session_id: The session ID to associate with the current context and its children. See https://insights.hanzo.ai/docs/data/sessions
Category:
Contexts
@@ -276,6 +286,39 @@ def get_context_distinct_id() -> Optional[str]:
return None
def set_context_device_id(device_id: str) -> None:
"""
Set the device ID for the current context, associating all feature flag requests in this or
child contexts with the given device ID (unless set_context_device_id is called again).
Entering a fresh context will clear the context-level device ID.
Args:
device_id: The device ID to associate with the current context and its children.
Category:
Contexts
"""
current_context = _get_current_context()
if current_context:
current_context.set_device_id(device_id)
def get_context_device_id() -> Optional[str]:
"""
Get the device ID for the current context.
Returns:
The device ID if set, None otherwise
Category:
Contexts
"""
current_context = _get_current_context()
if current_context:
return current_context.get_device_id()
return None
def set_capture_exception_code_variables_context(enabled: bool) -> None:
"""
Set whether code variables are captured for the current context.
@@ -330,20 +373,20 @@ F = TypeVar("F", bound=Callable[..., Any])
def scoped(fresh: bool = False, capture_exceptions: bool = True):
"""
Decorator that creates a new context for the function. Simply wraps
the function in a with posthog.new_context(): block.
the function in a with hanzo_insights.new_context(): block.
Args:
fresh: Whether to start with a fresh context (default: False)
capture_exceptions: Whether to capture and track exceptions with posthog error tracking (default: True)
capture_exceptions: Whether to capture and track exceptions with Insights error tracking (default: True)
Example:
@posthog.scoped()
@hanzo_insights.scoped()
def process_payment(payment_id):
posthog.tag("payment_id", payment_id)
posthog.tag("payment_method", "credit_card")
hanzo_insights.tag("payment_id", payment_id)
hanzo_insights.tag("payment_method", "credit_card")
# This event will be captured with tags
posthog.capture("payment_started")
hanzo_insights.capture("payment_started")
# If this raises an exception, it will be captured with tags
# and then re-raised
some_risky_function()
@@ -9,13 +9,13 @@ import threading
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from posthog.client import Client
from hanzo_insights.client import Client
class ExceptionCapture:
# TODO: Add client side rate limiting to prevent spamming the server with exceptions
log = logging.getLogger("posthog")
log = logging.getLogger("hanzo_insights")
def __init__(self, client: "Client"):
self.client = client
@@ -14,23 +14,23 @@ import types
from datetime import datetime
from types import FrameType, TracebackType # noqa: F401
from typing import ( # noqa: F401
TYPE_CHECKING,
Any,
Dict,
Iterator,
List,
Literal,
Optional,
Pattern,
Set,
Tuple,
TypedDict,
TypeVar,
Union,
cast,
TYPE_CHECKING,
Pattern,
)
from posthog.args import ExcInfo, ExceptionArg # noqa: F401
from hanzo_insights.args import ExceptionArg, ExcInfo # noqa: F401
try:
# Python 3.11
@@ -43,22 +43,31 @@ except ImportError:
DEFAULT_MAX_VALUE_LENGTH = 1024
DEFAULT_CODE_VARIABLES_MASK_PATTERNS = [
r"(?i).*password.*",
r"(?i).*secret.*",
r"(?i).*passwd.*",
r"(?i).*pwd.*",
r"(?i).*api_key.*",
r"(?i).*apikey.*",
r"(?i).*auth.*",
r"(?i).*credentials.*",
r"(?i).*privatekey.*",
r"(?i).*private_key.*",
r"(?i).*token.*",
r"(?i)password",
r"(?i)secret",
r"(?i)passwd",
r"(?i)pwd",
r"(?i)api_key",
r"(?i)apikey",
r"(?i)auth",
r"(?i)credentials",
r"(?i)privatekey",
r"(?i)private_key",
r"(?i)token",
r"(?i)aws_access_key_id",
r"(?i)_pass",
r"(?i)sk_",
r"(?i)jwt",
]
DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS = [r"^__.*"]
CODE_VARIABLES_REDACTED_VALUE = "$$_posthog_redacted_based_on_masking_rules_$$"
CODE_VARIABLES_REDACTED_VALUE = "$$_insights_redacted_based_on_masking_rules_$$"
CODE_VARIABLES_TOO_LONG_VALUE = "$$_insights_value_too_long_$$"
_MAX_VALUE_LENGTH_FOR_PATTERN_MATCH = 5_000
_MAX_COLLECTION_ITEMS_TO_SCAN = 100
_REGEX_METACHARACTERS = frozenset(r"\.^$*+?{}[]|()")
DEFAULT_TOTAL_VARIABLES_SIZE_LIMIT = 20 * 1024
@@ -759,12 +768,12 @@ def set_in_app_in_frames(frames, in_app_exclude, in_app_include, project_root=No
def exception_is_already_captured(error):
# type: (ExceptionArg) -> bool
if isinstance(error, BaseException):
return hasattr(error, "__posthog_exception_captured")
return hasattr(error, "__insights_exception_captured")
# Autocaptured exceptions are passed as a tuple from our system hooks,
# the second item is the exception value (the first is the exception type)
elif isinstance(error, tuple) and len(error) > 1:
return error[1] is not None and hasattr(
error[1], "__posthog_exception_captured"
error[1], "__insights_exception_captured"
)
else:
return False # type: ignore[unreachable]
@@ -773,14 +782,14 @@ def exception_is_already_captured(error):
def mark_exception_as_captured(error, uuid):
# type: (ExceptionArg, str) -> None
if isinstance(error, BaseException):
setattr(error, "__posthog_exception_captured", True)
setattr(error, "__posthog_exception_uuid", uuid)
setattr(error, "__insights_exception_captured", True)
setattr(error, "__insights_exception_uuid", uuid)
# Autocaptured exceptions are passed as a tuple from our system hooks,
# the second item is the exception value (the first is the exception type)
elif isinstance(error, tuple) and len(error) > 1:
if error[1] is not None:
setattr(error[1], "__posthog_exception_captured", True)
setattr(error[1], "__posthog_exception_uuid", uuid)
setattr(error[1], "__insights_exception_captured", True)
setattr(error[1], "__insights_exception_uuid", uuid)
def exc_info_from_error(error):
@@ -924,24 +933,95 @@ def strip_string(value, max_length=None):
)
def _extract_plain_substring(pattern):
# Matches inline flag groups like (?i), (?ai), (?ims), etc. that include the 'i' flag.
# Python regex flags: a=ASCII, i=IGNORECASE, L=LOCALE, m=MULTILINE, s=DOTALL, u=UNICODE, x=VERBOSE
inline_flags = re.match(r"^\(\?[aiLmsux]*i[aiLmsux]*\)", pattern)
if not inline_flags:
return None
remainder = pattern[inline_flags.end() :]
if not remainder or any(c in _REGEX_METACHARACTERS for c in remainder):
return None
return remainder.lower()
def _compile_patterns(patterns):
compiled = []
if not patterns:
return None
substrings = []
regexes = []
for pattern in patterns:
try:
compiled.append(re.compile(pattern))
except:
pass
return compiled
simple = _extract_plain_substring(pattern)
if simple is not None:
substrings.append(simple)
else:
try:
regexes.append(re.compile(pattern))
except Exception:
pass
if not substrings and not regexes:
return None
return (substrings, regexes)
def _pattern_matches(name, patterns):
for pattern in patterns:
if patterns is None:
return False
substrings, regexes = patterns
if substrings:
name_lower = name.lower()
for s in substrings:
if s in name_lower:
return True
for pattern in regexes:
if pattern.search(name):
return True
return False
def _serialize_variable_value(value, limiter, max_length=1024):
def _mask_sensitive_data(value, compiled_mask, _seen=None):
if not compiled_mask:
return value
if isinstance(value, (dict, list, tuple)):
if _seen is None:
_seen = set()
obj_id = id(value)
if obj_id in _seen:
return "<circular ref>"
_seen.add(obj_id)
if isinstance(value, dict):
if len(value) > _MAX_COLLECTION_ITEMS_TO_SCAN:
return CODE_VARIABLES_TOO_LONG_VALUE
result = {}
for k, v in value.items():
key_str = str(k) if not isinstance(k, str) else k
if len(key_str) > _MAX_VALUE_LENGTH_FOR_PATTERN_MATCH:
result[k] = CODE_VARIABLES_TOO_LONG_VALUE
elif _pattern_matches(key_str, compiled_mask):
result[k] = CODE_VARIABLES_REDACTED_VALUE
else:
result[k] = _mask_sensitive_data(v, compiled_mask, _seen)
return result
elif isinstance(value, (list, tuple)):
if len(value) > _MAX_COLLECTION_ITEMS_TO_SCAN:
return CODE_VARIABLES_TOO_LONG_VALUE
masked_items = [
_mask_sensitive_data(item, compiled_mask, _seen) for item in value
]
return type(value)(masked_items)
elif isinstance(value, str):
if len(value) > _MAX_VALUE_LENGTH_FOR_PATTERN_MATCH:
return CODE_VARIABLES_TOO_LONG_VALUE
if _pattern_matches(value, compiled_mask):
return CODE_VARIABLES_REDACTED_VALUE
return value
else:
return value
def _serialize_variable_value(value, limiter, max_length=1024, compiled_mask=None):
try:
if value is None:
result = "None"
@@ -954,9 +1034,15 @@ def _serialize_variable_value(value, limiter, max_length=1024):
limiter.add(result_size)
return value
elif isinstance(value, str):
result = value
if len(value) > _MAX_VALUE_LENGTH_FOR_PATTERN_MATCH:
result = CODE_VARIABLES_TOO_LONG_VALUE
elif compiled_mask and _pattern_matches(value, compiled_mask):
result = CODE_VARIABLES_REDACTED_VALUE
else:
result = value
else:
result = json.dumps(value)
masked_value = _mask_sensitive_data(value, compiled_mask)
result = json.dumps(masked_value)
if len(result) > max_length:
result = result[: max_length - 3] + "..."
@@ -969,19 +1055,30 @@ def _serialize_variable_value(value, limiter, max_length=1024):
return result
except Exception:
try:
fallback = f"<{type(value).__name__}>"
fallback_size = len(fallback)
if not limiter.can_add(fallback_size):
result = repr(value)
if len(result) > max_length:
result = result[: max_length - 3] + "..."
result_size = len(result)
if not limiter.can_add(result_size):
return None
limiter.add(fallback_size)
return fallback
limiter.add(result_size)
return result
except Exception:
fallback = "<unserializable object>"
fallback_size = len(fallback)
if not limiter.can_add(fallback_size):
return None
limiter.add(fallback_size)
return fallback
try:
fallback = f"<{type(value).__name__}>"
fallback_size = len(fallback)
if not limiter.can_add(fallback_size):
return None
limiter.add(fallback_size)
return fallback
except Exception:
fallback = "<unserializable object>"
fallback_size = len(fallback)
if not limiter.can_add(fallback_size):
return None
limiter.add(fallback_size)
return fallback
def _is_simple_type(value):
@@ -1032,7 +1129,9 @@ def serialize_code_variables(
limiter.add(redacted_size)
result[name] = redacted_value
else:
serialized = _serialize_variable_value(value, limiter, max_length)
serialized = _serialize_variable_value(
value, limiter, max_length, compiled_mask
)
if serialized is None:
break
result[name] = serialized
@@ -1042,6 +1141,17 @@ def serialize_code_variables(
def try_attach_code_variables_to_frames(
all_exceptions, exc_info, mask_patterns, ignore_patterns
):
try:
attach_code_variables_to_frames(
all_exceptions, exc_info, mask_patterns, ignore_patterns
)
except Exception:
pass
def attach_code_variables_to_frames(
all_exceptions, exc_info, mask_patterns, ignore_patterns
):
exc_type, exc_value, traceback = exc_info
@@ -2,21 +2,46 @@ import datetime
import hashlib
import logging
import re
import warnings
from typing import Optional
from dateutil import parser
from dateutil.relativedelta import relativedelta
from posthog import utils
from posthog.types import FlagValue
from posthog.utils import convert_to_datetime_aware, is_valid_regex
from hanzo_insights import utils
from hanzo_insights.types import FlagValue
from hanzo_insights.utils import convert_to_datetime_aware, is_valid_regex
__LONG_SCALE__ = float(0xFFFFFFFFFFFFFFF)
log = logging.getLogger("posthog")
log = logging.getLogger("hanzo_insights")
NONE_VALUES_ALLOWED_OPERATORS = ["is_not"]
# All operators supported by match_property, grouped by category.
EQUALITY_OPERATORS = ("exact", "is_not", "is_set", "is_not_set")
STRING_OPERATORS = ("icontains", "not_icontains", "regex", "not_regex")
NUMERIC_OPERATORS = ("gt", "gte", "lt", "lte")
DATE_OPERATORS = ("is_date_before", "is_date_after")
SEMVER_COMPARISON_OPERATORS = (
"semver_eq",
"semver_neq",
"semver_gt",
"semver_gte",
"semver_lt",
"semver_lte",
)
SEMVER_RANGE_OPERATORS = ("semver_tilde", "semver_caret", "semver_wildcard")
SEMVER_OPERATORS = SEMVER_COMPARISON_OPERATORS + SEMVER_RANGE_OPERATORS
PROPERTY_OPERATORS = (
EQUALITY_OPERATORS
+ STRING_OPERATORS
+ NUMERIC_OPERATORS
+ DATE_OPERATORS
+ SEMVER_OPERATORS
)
class InconclusiveMatchError(Exception):
pass
@@ -34,18 +59,18 @@ class RequiresServerEvaluation(Exception):
pass
# This function takes a distinct_id and a feature flag key and returns a float between 0 and 1.
# Given the same distinct_id and key, it'll always return the same float. These floats are
# This function takes a bucketing value and a feature flag key and returns a float between 0 and 1.
# Given the same bucketing value and key, it'll always return the same float. These floats are
# uniformly distributed between 0 and 1, so if we want to show this feature to 20% of traffic
# we can do _hash(key, distinct_id) < 0.2
def _hash(key: str, distinct_id: str, salt: str = "") -> float:
hash_key = f"{key}.{distinct_id}{salt}"
# we can do _hash(key, bucketing_value) < 0.2
def _hash(key: str, bucketing_value: str, salt: str = "") -> float:
hash_key = f"{key}.{bucketing_value}{salt}"
hash_val = int(hashlib.sha1(hash_key.encode("utf-8")).hexdigest()[:15], 16)
return hash_val / __LONG_SCALE__
def get_matching_variant(flag, distinct_id):
hash_value = _hash(flag["key"], distinct_id, salt="variant")
def get_matching_variant(flag, bucketing_value):
hash_value = _hash(flag["key"], bucketing_value, salt="variant")
for variant in variant_lookup_table(flag):
if hash_value >= variant["value_min"] and hash_value < variant["value_max"]:
return variant["key"]
@@ -68,7 +93,13 @@ def variant_lookup_table(feature_flag):
def evaluate_flag_dependency(
property, flags_by_key, evaluation_cache, distinct_id, properties, cohort_properties
property,
flags_by_key,
evaluation_cache,
distinct_id,
properties,
cohort_properties,
device_id=None,
):
"""
Evaluate a flag dependency property according to the dependency chain algorithm.
@@ -80,6 +111,7 @@ def evaluate_flag_dependency(
distinct_id: The distinct ID being evaluated
properties: Person properties for evaluation
cohort_properties: Cohort properties for evaluation
device_id: The device ID for bucketing (optional)
Returns:
bool: True if all dependencies in the chain evaluate to True, False otherwise
@@ -124,13 +156,27 @@ def evaluate_flag_dependency(
else:
# Recursively evaluate the dependency
try:
dep_flag_filters = dep_flag.get("filters") or {}
dep_aggregation_group_type_index = dep_flag_filters.get(
"aggregation_group_type_index"
)
if dep_aggregation_group_type_index is not None:
# Group flags should continue bucketing by the group key
# from the current evaluation context.
dep_bucketing_value = distinct_id
else:
dep_bucketing_value = resolve_bucketing_value(
dep_flag, distinct_id, device_id
)
dep_result = match_feature_flag_properties(
dep_flag,
distinct_id,
properties,
cohort_properties,
flags_by_key,
evaluation_cache,
cohort_properties=cohort_properties,
flags_by_key=flags_by_key,
evaluation_cache=evaluation_cache,
device_id=device_id,
bucketing_value=dep_bucketing_value,
)
evaluation_cache[dep_flag_key] = dep_result
except InconclusiveMatchError as e:
@@ -215,21 +261,54 @@ def matches_dependency_value(expected_value, actual_value):
return False
def resolve_bucketing_value(flag, distinct_id, device_id=None):
"""Resolve the bucketing value for a flag based on its bucketing_identifier setting.
Returns:
The appropriate identifier string to use for hashing/bucketing.
Raises:
InconclusiveMatchError: If the flag requires device_id but none was provided.
"""
flag_filters = flag.get("filters") or {}
bucketing_identifier = flag.get("bucketing_identifier") or flag_filters.get(
"bucketing_identifier"
)
if bucketing_identifier == "device_id":
if not device_id:
raise InconclusiveMatchError(
"Flag requires device_id for bucketing but none was provided"
)
return device_id
return distinct_id
def match_feature_flag_properties(
flag,
distinct_id,
properties,
*,
cohort_properties=None,
flags_by_key=None,
evaluation_cache=None,
device_id=None,
bucketing_value=None,
) -> FlagValue:
flag_conditions = (flag.get("filters") or {}).get("groups") or []
if bucketing_value is None:
warnings.warn(
"Calling match_feature_flag_properties() without bucketing_value is deprecated. "
"Pass bucketing_value explicitly. This fallback will be removed in a future major release.",
DeprecationWarning,
stacklevel=2,
)
bucketing_value = resolve_bucketing_value(flag, distinct_id, device_id)
flag_filters = flag.get("filters") or {}
flag_conditions = flag_filters.get("groups") or []
is_inconclusive = False
cohort_properties = cohort_properties or {}
# Some filters can be explicitly set to null, which require accessing variants like so
flag_variants = ((flag.get("filters") or {}).get("multivariate") or {}).get(
"variants"
) or []
flag_variants = (flag_filters.get("multivariate") or {}).get("variants") or []
valid_variant_keys = [variant["key"] for variant in flag_variants]
for condition in flag_conditions:
@@ -244,12 +323,14 @@ def match_feature_flag_properties(
cohort_properties,
flags_by_key,
evaluation_cache,
bucketing_value=bucketing_value,
device_id=device_id,
):
variant_override = condition.get("variant")
if variant_override and variant_override in valid_variant_keys:
variant = variant_override
else:
variant = get_matching_variant(flag, distinct_id)
variant = get_matching_variant(flag, bucketing_value)
return variant or True
except RequiresServerEvaluation:
# Static cohort or other missing server-side data - must fallback to API
@@ -277,6 +358,9 @@ def is_condition_match(
cohort_properties,
flags_by_key=None,
evaluation_cache=None,
*,
bucketing_value,
device_id=None,
) -> bool:
rollout_percentage = condition.get("rollout_percentage")
if len(condition.get("properties") or []) > 0:
@@ -290,6 +374,7 @@ def is_condition_match(
flags_by_key,
evaluation_cache,
distinct_id,
device_id=device_id,
)
elif property_type == "flag":
matches = evaluate_flag_dependency(
@@ -299,6 +384,7 @@ def is_condition_match(
distinct_id,
properties,
cohort_properties,
device_id=device_id,
)
else:
matches = match_property(prop, properties)
@@ -308,9 +394,9 @@ def is_condition_match(
if rollout_percentage is None:
return True
if rollout_percentage is not None and _hash(feature_flag["key"], distinct_id) > (
rollout_percentage / 100
):
if rollout_percentage is not None and _hash(
feature_flag["key"], bucketing_value
) > (rollout_percentage / 100):
return False
return True
@@ -323,6 +409,9 @@ def match_property(property, property_values) -> bool:
operator = property.get("operator") or "exact"
value = property.get("value")
if operator not in PROPERTY_OPERATORS:
raise InconclusiveMatchError(f"Unknown operator {operator}")
if key not in property_values:
raise InconclusiveMatchError(
"can't match properties without a given property value"
@@ -443,7 +532,64 @@ def match_property(property, property_values) -> bool:
"The date provided must be a string or date object"
)
# if we get here, we don't know how to handle the operator
if operator in SEMVER_OPERATORS:
try:
override_parsed = parse_semver(override_value)
except (ValueError, TypeError):
raise InconclusiveMatchError(
f"Person property value '{override_value}' is not a valid semver"
)
if operator in SEMVER_COMPARISON_OPERATORS:
try:
flag_parsed = parse_semver(value)
except (ValueError, TypeError):
raise InconclusiveMatchError(
f"Flag semver value '{value}' is not a valid semver"
)
if operator == "semver_eq":
return override_parsed == flag_parsed
elif operator == "semver_neq":
return override_parsed != flag_parsed
elif operator == "semver_gt":
return override_parsed > flag_parsed
elif operator == "semver_gte":
return override_parsed >= flag_parsed
elif operator == "semver_lt":
return override_parsed < flag_parsed
elif operator == "semver_lte":
return override_parsed <= flag_parsed
elif operator == "semver_tilde":
try:
lower, upper = _tilde_bounds(str(value))
except (ValueError, TypeError):
raise InconclusiveMatchError(
f"Flag semver value '{value}' is not valid for tilde operator"
)
return lower <= override_parsed < upper
elif operator == "semver_caret":
try:
lower, upper = _caret_bounds(str(value))
except (ValueError, TypeError):
raise InconclusiveMatchError(
f"Flag semver value '{value}' is not valid for caret operator"
)
return lower <= override_parsed < upper
elif operator == "semver_wildcard":
try:
lower, upper = _wildcard_bounds(str(value))
except (ValueError, TypeError):
raise InconclusiveMatchError(
f"Flag semver value '{value}' is not valid for wildcard operator"
)
return lower <= override_parsed < upper
# Unreachable: all operators in PROPERTY_OPERATORS are handled above,
# and unknown operators are rejected at the top of this function.
raise InconclusiveMatchError(f"Unknown operator {operator}")
@@ -454,6 +600,7 @@ def match_cohort(
flags_by_key=None,
evaluation_cache=None,
distinct_id=None,
device_id=None,
) -> bool:
# Cohort properties are in the form of property groups like this:
# {
@@ -478,6 +625,7 @@ def match_cohort(
flags_by_key,
evaluation_cache,
distinct_id,
device_id=device_id,
)
@@ -488,6 +636,7 @@ def match_property_group(
flags_by_key=None,
evaluation_cache=None,
distinct_id=None,
device_id=None,
) -> bool:
if not property_group:
return True
@@ -512,6 +661,7 @@ def match_property_group(
flags_by_key,
evaluation_cache,
distinct_id,
device_id=device_id,
)
if property_group_type == "AND":
if not matches:
@@ -545,6 +695,7 @@ def match_property_group(
flags_by_key,
evaluation_cache,
distinct_id,
device_id=device_id,
)
elif prop.get("type") == "flag":
matches = evaluate_flag_dependency(
@@ -554,6 +705,7 @@ def match_property_group(
distinct_id,
property_values,
cohort_properties,
device_id=device_id,
)
else:
matches = match_property(prop, property_values)
@@ -618,3 +770,75 @@ def relative_date_parse_for_feature_flag_matching(
return parsed_dt
else:
return None
def parse_semver(value: str) -> tuple:
"""Parse a semver string into a comparable (major, minor, patch) integer tuple.
Matches the behavior of the sortableSemver HogQL function:
- Handles v-prefix, whitespace, pre-release suffixes
- Defaults missing components to 0 (e.g., 1.2 -> 1.2.0)
Raises ValueError if parsing fails.
"""
text = str(value).strip().lstrip("vV")
# Strip pre-release/build metadata suffix
text = text.split("-")[0].split("+")[0]
parts = text.split(".")
if not parts or not parts[0]:
raise ValueError("Invalid semver format")
major = int(parts[0])
minor = int(parts[1]) if len(parts) > 1 and parts[1] else 0
patch = int(parts[2]) if len(parts) > 2 and parts[2] else 0
return (major, minor, patch)
def _tilde_bounds(value: str) -> tuple:
"""~1.2.3 means >=1.2.3 <1.3.0 (allows patch-level changes)."""
major, minor, patch = parse_semver(value)
return (major, minor, patch), (major, minor + 1, 0)
def _caret_bounds(value: str) -> tuple:
"""Caret follows semver spec:
^1.2.3 means >=1.2.3 <2.0.0
^0.2.3 means >=0.2.3 <0.3.0
^0.0.3 means >=0.0.3 <0.0.4
"""
major, minor, patch = parse_semver(value)
lower = (major, minor, patch)
if major > 0:
upper = (major + 1, 0, 0)
elif minor > 0:
upper = (0, minor + 1, 0)
else:
upper = (0, 0, patch + 1)
return lower, upper
def _wildcard_bounds(value: str) -> tuple:
"""Wildcard matching:
1.* means >=1.0.0 <2.0.0
1.2.* means >=1.2.0 <1.3.0
"""
cleaned = str(value).strip().lstrip("vV").replace("*", "").rstrip(".")
if not cleaned:
raise ValueError("Invalid wildcard pattern")
parts = [p for p in cleaned.split(".") if p]
if not parts:
raise ValueError("Invalid wildcard pattern")
if len(parts) == 1:
major = int(parts[0])
return (major, 0, 0), (major + 1, 0, 0)
elif len(parts) == 2:
major, minor = int(parts[0]), int(parts[1])
return (major, minor, 0), (major, minor + 1, 0)
else:
major, minor, patch = int(parts[0]), int(parts[1]), int(parts[2])
return (major, minor, patch), (major, minor, patch + 1)
+127
View File
@@ -0,0 +1,127 @@
"""
Flag Definition Cache Provider interface for multi-worker environments.
EXPERIMENTAL: This API may change in future minor version bumps.
This module provides an interface for external caching of feature flag definitions,
enabling multi-worker environments (Kubernetes, load-balanced servers, serverless
functions) to share flag definitions and reduce API calls.
Usage:
from hanzo_insights import Insights
from hanzo_insights.flag_definition_cache import FlagDefinitionCacheProvider
cache = RedisFlagDefinitionCache(redis_client, "my-team")
client = Insights(
"<project_api_key>",
personal_api_key="<personal_api_key>",
flag_definition_cache_provider=cache,
)
"""
from typing import Any, Dict, List, Optional, Protocol, runtime_checkable
from typing_extensions import Required, TypedDict
class FlagDefinitionCacheData(TypedDict):
"""
Data structure for cached flag definitions.
Attributes:
flags: List of feature flag definition dictionaries from the API.
group_type_mapping: Mapping of group type indices to group names.
cohorts: Dictionary of cohort definitions for local evaluation.
"""
flags: Required[List[Dict[str, Any]]]
group_type_mapping: Required[Dict[str, str]]
cohorts: Required[Dict[str, Any]]
@runtime_checkable
class FlagDefinitionCacheProvider(Protocol):
"""
Interface for external caching of feature flag definitions.
Enables multi-worker environments to share flag definitions, reducing API
calls while ensuring all workers have consistent data.
EXPERIMENTAL: This API may change in future minor version bumps.
The four methods handle the complete lifecycle of flag definition caching:
1. `should_fetch_flag_definitions()` - Called before each poll to determine
if this worker should fetch new definitions. Use for distributed lock
coordination to ensure only one worker fetches at a time.
2. `get_flag_definitions()` - Called when `should_fetch_flag_definitions()`
returns False. Returns cached definitions if available.
3. `on_flag_definitions_received()` - Called after successfully fetching
new definitions from the API. Store the data in your external cache
and release any locks.
4. `shutdown()` - Called when the Insights client shuts down. Release any
distributed locks and clean up resources.
Error Handling:
All methods are wrapped in try/except. Errors will be logged but will
never break flag evaluation. On error:
- `should_fetch_flag_definitions()` errors default to fetching (fail-safe)
- `get_flag_definitions()` errors fall back to API fetch
- `on_flag_definitions_received()` errors are logged but flags remain in memory
- `shutdown()` errors are logged but shutdown continues
"""
def get_flag_definitions(self) -> Optional[FlagDefinitionCacheData]:
"""
Retrieve cached flag definitions.
Returns:
Cached flag definitions if available and valid, None otherwise.
Returning None will trigger a fetch from the API if this worker
has no flags loaded yet.
"""
...
def should_fetch_flag_definitions(self) -> bool:
"""
Determine whether this instance should fetch new flag definitions.
Use this for distributed lock coordination. Only one worker should
return True to avoid thundering herd problems. A typical implementation
uses a distributed lock (e.g., Redis SETNX) that expires after the
poll interval.
Returns:
True if this instance should fetch from the API, False otherwise.
When False, the client will call `get_flag_definitions()` to
retrieve cached data instead.
"""
...
def on_flag_definitions_received(self, data: FlagDefinitionCacheData) -> None:
"""
Called after successfully receiving new flag definitions from Insights.
Use this to store the data in your external cache and release any
distributed locks acquired in `should_fetch_flag_definitions()`.
Args:
data: The flag definitions to cache, containing flags,
group_type_mapping, and cohorts.
"""
...
def shutdown(self) -> None:
"""
Called when the Insights client shuts down.
Use this to release any distributed locks and clean up resources.
This method is called even if `should_fetch_flag_definitions()`
returned False, so implementations should handle the case where
no lock was acquired.
"""
...
@@ -1,6 +1,6 @@
from typing import TYPE_CHECKING, cast
from posthog import contexts
from posthog.client import Client
from hanzo_insights import contexts
from hanzo_insights.client import Client
try:
from asgiref.sync import iscoroutinefunction, markcoroutinefunction
@@ -21,25 +21,26 @@ if TYPE_CHECKING:
from typing import Callable, Dict, Any, Optional, Union, Awaitable # noqa: F401
class PosthogContextMiddleware:
class InsightsContextMiddleware:
"""Middleware to automatically track Django requests.
This middleware wraps all calls with a posthog context. It attempts to extract the following from the request headers:
- Session ID, (extracted from `X-POSTHOG-SESSION-ID`)
- Distinct ID, (extracted from `X-POSTHOG-DISTINCT-ID`)
This middleware wraps all calls with an Insights context. It attempts to extract the following from the request headers:
- Session ID, (extracted from `X-INSIGHTS-SESSION-ID`)
- Distinct ID, (extracted from `X-INSIGHTS-DISTINCT-ID`)
- Request URL as $current_url
- Request Method as $request_method
The context will also auto-capture exceptions and send them to PostHog, unless you disable it by setting
`POSTHOG_MW_CAPTURE_EXCEPTIONS` to `False` in your Django settings. The exceptions are captured using the
global client, unless the setting `POSTHOG_MW_CLIENT` is set to a custom client instance
The context will also auto-capture exceptions and send them to Insights, unless you disable it by setting
`INSIGHTS_MW_CAPTURE_EXCEPTIONS` to `False` in your Django settings.
The exceptions are captured using the global client, unless the setting `INSIGHTS_MW_CLIENT`
is set to a custom client instance.
The middleware behaviour is customisable through 3 additional functions:
- `POSTHOG_MW_EXTRA_TAGS`, which is a Callable[[HttpRequest], Dict[str, Any]] expected to return a dictionary of additional tags to be added to the context.
- `POSTHOG_MW_REQUEST_FILTER`, which is a Callable[[HttpRequest], bool] expected to return `False` if the request should not be tracked.
- `POSTHOG_MW_TAG_MAP`, which is a Callable[[Dict[str, Any]], Dict[str, Any]], which you can use to modify the tags before they're added to the context.
- `INSIGHTS_MW_EXTRA_TAGS`, which is a Callable[[HttpRequest], Dict[str, Any]] expected to return a dictionary of additional tags to be added to the context.
- `INSIGHTS_MW_REQUEST_FILTER`, which is a Callable[[HttpRequest], bool] expected to return `False` if the request should not be tracked.
- `INSIGHTS_MW_TAG_MAP`, which is a Callable[[Dict[str, Any]], Dict[str, Any]], which you can use to modify the tags before they're added to the context.
You can use the `POSTHOG_MW_TAG_MAP` function to remove any default tags you don't want to capture, or override them with your own values.
You can use the `INSIGHTS_MW_TAG_MAP` function to remove any default tags you don't want to capture, or override them with your own values.
Context tags are automatically included as properties on all events captured within a context, including exceptions.
See the context documentation for more information. The extracted distinct ID and session ID, if found, are used to
@@ -66,47 +67,48 @@ class PosthogContextMiddleware:
from django.conf import settings
if hasattr(settings, "POSTHOG_MW_EXTRA_TAGS") and callable(
settings.POSTHOG_MW_EXTRA_TAGS
):
def _get_setting(name):
insights_name = f"INSIGHTS_MW_{name}"
if hasattr(settings, insights_name):
return getattr(settings, insights_name)
return None
extra_tags = _get_setting("EXTRA_TAGS")
if extra_tags and callable(extra_tags):
self.extra_tags = cast(
"Optional[Callable[[HttpRequest], Dict[str, Any]]]",
settings.POSTHOG_MW_EXTRA_TAGS,
extra_tags,
)
else:
self.extra_tags = None
if hasattr(settings, "POSTHOG_MW_REQUEST_FILTER") and callable(
settings.POSTHOG_MW_REQUEST_FILTER
):
request_filter = _get_setting("REQUEST_FILTER")
if request_filter and callable(request_filter):
self.request_filter = cast(
"Optional[Callable[[HttpRequest], bool]]",
settings.POSTHOG_MW_REQUEST_FILTER,
request_filter,
)
else:
self.request_filter = None
if hasattr(settings, "POSTHOG_MW_TAG_MAP") and callable(
settings.POSTHOG_MW_TAG_MAP
):
tag_map = _get_setting("TAG_MAP")
if tag_map and callable(tag_map):
self.tag_map = cast(
"Optional[Callable[[Dict[str, Any]], Dict[str, Any]]]",
settings.POSTHOG_MW_TAG_MAP,
tag_map,
)
else:
self.tag_map = None
if hasattr(settings, "POSTHOG_MW_CAPTURE_EXCEPTIONS") and isinstance(
settings.POSTHOG_MW_CAPTURE_EXCEPTIONS, bool
):
self.capture_exceptions = settings.POSTHOG_MW_CAPTURE_EXCEPTIONS
capture_exceptions = _get_setting("CAPTURE_EXCEPTIONS")
if isinstance(capture_exceptions, bool):
self.capture_exceptions = capture_exceptions
else:
self.capture_exceptions = True
if hasattr(settings, "POSTHOG_MW_CLIENT") and isinstance(
settings.POSTHOG_MW_CLIENT, Client
):
self.client = cast("Optional[Client]", settings.POSTHOG_MW_CLIENT)
mw_client = _get_setting("CLIENT")
if isinstance(mw_client, Client):
self.client = cast("Optional[Client]", mw_client)
else:
self.client = None
@@ -125,13 +127,13 @@ class PosthogContextMiddleware:
"""
tags = {}
# Extract session ID from X-POSTHOG-SESSION-ID header
session_id = request.headers.get("X-POSTHOG-SESSION-ID")
# Extract session ID from X-INSIGHTS-SESSION-ID header
session_id = request.headers.get("X-INSIGHTS-SESSION-ID")
if session_id:
contexts.set_context_session(session_id)
# Extract distinct ID from X-POSTHOG-DISTINCT-ID header or request user id
distinct_id = request.headers.get("X-POSTHOG-DISTINCT-ID") or user_id
# Extract distinct ID from X-INSIGHTS-DISTINCT-ID header or request user id
distinct_id = request.headers.get("X-INSIGHTS-DISTINCT-ID") or user_id
if distinct_id:
contexts.identify_context(distinct_id)
@@ -155,7 +157,7 @@ class PosthogContextMiddleware:
# Extract IP address
ip_address = request.headers.get("X-Forwarded-For")
if ip_address:
tags["$ip_address"] = ip_address
tags["$ip"] = ip_address
# Extract user agent
user_agent = request.headers.get("User-Agent")
@@ -314,6 +316,6 @@ class PosthogContextMiddleware:
if self.client:
self.client.capture_exception(exception)
else:
from posthog import capture_exception
from hanzo_insights import capture_exception
capture_exception(exception)
+397
View File
@@ -0,0 +1,397 @@
import json
import logging
import re
import socket
from dataclasses import dataclass
from datetime import date, datetime, timezone
from gzip import GzipFile
from io import BytesIO
from typing import Any, List, Optional, Tuple, Union
import requests
from dateutil.tz import tzutc
from requests.adapters import HTTPAdapter # type: ignore[import-untyped]
from urllib3.connection import HTTPConnection
from urllib3.util.retry import Retry
from hanzo_insights.utils import remove_trailing_slash
from hanzo_insights.version import VERSION
SocketOptions = List[Tuple[int, int, Union[int, bytes]]]
KEEPALIVE_IDLE_SECONDS = 60
KEEPALIVE_INTERVAL_SECONDS = 60
KEEPALIVE_PROBE_COUNT = 3
# TCP keepalive probes idle connections to prevent them from being dropped.
# SO_KEEPALIVE is cross-platform, but timing options vary:
# - Linux: TCP_KEEPIDLE, TCP_KEEPINTVL, TCP_KEEPCNT
# - macOS: only SO_KEEPALIVE (uses system defaults)
# - Windows: TCP_KEEPIDLE, TCP_KEEPINTVL (since Windows 10 1709)
KEEP_ALIVE_SOCKET_OPTIONS: SocketOptions = list(
HTTPConnection.default_socket_options
) + [
(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1),
]
for attr, value in [
("TCP_KEEPIDLE", KEEPALIVE_IDLE_SECONDS),
("TCP_KEEPINTVL", KEEPALIVE_INTERVAL_SECONDS),
("TCP_KEEPCNT", KEEPALIVE_PROBE_COUNT),
]:
if hasattr(socket, attr):
KEEP_ALIVE_SOCKET_OPTIONS.append((socket.SOL_TCP, getattr(socket, attr), value))
# Status codes that indicate transient server errors worth retrying
RETRY_STATUS_FORCELIST = [408, 500, 502, 503, 504]
def _mask_tokens_in_url(url: str) -> str:
"""Mask token values in URLs for safe logging, keeping first 10 chars visible."""
return re.sub(r"(token=)([^&]{10})[^&]*", r"\1\2...", url)
@dataclass
class GetResponse:
"""Response from a GET request with ETag support."""
data: Any
etag: Optional[str] = None
not_modified: bool = False
class HTTPAdapterWithSocketOptions(HTTPAdapter):
"""HTTPAdapter with configurable socket options."""
def __init__(self, *args, socket_options: Optional[SocketOptions] = None, **kwargs):
self.socket_options = socket_options
super().__init__(*args, **kwargs)
def init_poolmanager(self, *args, **kwargs):
if self.socket_options is not None:
kwargs["socket_options"] = self.socket_options
super().init_poolmanager(*args, **kwargs)
def _build_session(socket_options: Optional[SocketOptions] = None) -> requests.Session:
"""Build a session for general requests (batch, decide, etc.)."""
adapter = HTTPAdapterWithSocketOptions(
max_retries=Retry(
total=2,
connect=2,
read=2,
),
socket_options=socket_options,
)
session = requests.Session()
session.mount("https://", adapter)
return session
def _build_flags_session(
socket_options: Optional[SocketOptions] = None,
) -> requests.Session:
"""
Build a session for feature flag requests with POST retries.
Feature flag requests are idempotent (read-only), so retrying POST
requests is safe. This session retries on transient server errors
(408, 5xx) and network failures with exponential backoff
(0.5s, 1s delays between retries).
"""
adapter = HTTPAdapterWithSocketOptions(
max_retries=Retry(
total=2,
connect=2,
read=2,
backoff_factor=0.5,
status_forcelist=RETRY_STATUS_FORCELIST,
allowed_methods=["POST"],
),
socket_options=socket_options,
)
session = requests.Session()
session.mount("https://", adapter)
return session
_session = _build_session()
_flags_session = _build_flags_session()
_socket_options: Optional[SocketOptions] = None
_pooling_enabled = True
def _get_session() -> requests.Session:
if _pooling_enabled:
return _session
return _build_session(_socket_options)
def _get_flags_session() -> requests.Session:
if _pooling_enabled:
return _flags_session
return _build_flags_session(_socket_options)
def set_socket_options(socket_options: Optional[SocketOptions]) -> None:
"""
Configure socket options for all HTTP connections.
Example:
from hanzo_insights import set_socket_options
set_socket_options([(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)])
"""
global _session, _flags_session, _socket_options
if socket_options == _socket_options:
return
_socket_options = socket_options
_session = _build_session(socket_options)
_flags_session = _build_flags_session(socket_options)
def enable_keep_alive() -> None:
"""Enable TCP keepalive to prevent idle connections from being dropped."""
set_socket_options(KEEP_ALIVE_SOCKET_OPTIONS)
def disable_connection_reuse() -> None:
"""Disable connection reuse, creating a fresh connection for each request."""
global _pooling_enabled
_pooling_enabled = False
US_INGESTION_ENDPOINT = "https://us.i.insights.hanzo.ai"
EU_INGESTION_ENDPOINT = "https://eu.i.insights.hanzo.ai"
DEFAULT_HOST = US_INGESTION_ENDPOINT
USER_AGENT = "hanzo-insights-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", "https://insights.hanzo.ai", "https://us.insights.hanzo.ai"):
return US_INGESTION_ENDPOINT
elif trimmed_host in ("https://eu.posthog.com", "https://eu.insights.hanzo.ai"):
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,
session: Optional[requests.Session] = None,
**kwargs,
) -> requests.Response:
"""Post the `kwargs` to the API"""
log = logging.getLogger("hanzo_insights")
body = kwargs
body["sentAt"] = datetime.now(tz=tzutc()).isoformat()
url = remove_trailing_slash(host or DEFAULT_HOST) + path
body["api_key"] = api_key
data = json.dumps(body, cls=DatetimeSerializer)
log.debug("making request: %s to url: %s", data, url)
headers = {"Content-Type": "application/json", "User-Agent": USER_AGENT}
if gzip:
headers["Content-Encoding"] = "gzip"
buf = BytesIO()
with GzipFile(fileobj=buf, mode="w") as gz:
# 'data' was produced by json.dumps(),
# whose default encoding is utf-8.
gz.write(data.encode("utf-8"))
data = buf.getvalue()
res = (session or _get_session()).post(
url, data=data, headers=headers, timeout=timeout
)
if res.status_code == 200:
log.debug("data uploaded successfully")
return res
def _process_response(
res: requests.Response, success_message: str, *, return_json: bool = True
) -> Union[requests.Response, Any]:
log = logging.getLogger("hanzo_insights")
if res.status_code == 200:
log.debug(success_message)
response = res.json() if return_json else res
# Handle quota limited decide responses by raising a specific error
# NB: other services also put entries into the quotaLimited key, but right now we only care about feature flags
# since most of the other services handle quota limiting in other places in the application.
if (
isinstance(response, dict)
and "quotaLimited" in response
and isinstance(response["quotaLimited"], list)
and "feature_flags" in response["quotaLimited"]
):
log.warning(
"[FEATURE FLAGS] Feature flags quota limited, resetting feature flag data. Learn more about billing limits at https://insights.hanzo.ai/docs/billing/limits-alerts"
)
raise QuotaLimitError(res.status_code, "Feature flags quota limited")
return response
retry_after = None
retry_after_header = res.headers.get("Retry-After")
if retry_after_header:
try:
retry_after = float(retry_after_header)
except (ValueError, TypeError):
try:
from email.utils import parsedate_to_datetime
retry_after = max(
0.0,
(
parsedate_to_datetime(retry_after_header)
- datetime.now(timezone.utc)
).total_seconds(),
)
except (ValueError, TypeError):
pass
try:
payload = res.json()
log.debug("received response: %s", payload)
raise APIError(res.status_code, payload["detail"], retry_after=retry_after)
except (KeyError, ValueError):
raise APIError(res.status_code, res.text, retry_after=retry_after)
def decide(
api_key: str,
host: Optional[str] = None,
gzip: bool = False,
timeout: int = 15,
**kwargs,
) -> Any:
"""Post the `kwargs to the decide API endpoint"""
res = post(api_key, host, "/decide/?v=4", gzip, timeout, **kwargs)
return _process_response(res, success_message="Feature flags decided successfully")
def flags(
api_key: str,
host: Optional[str] = None,
gzip: bool = False,
timeout: int = 15,
**kwargs,
) -> Any:
"""Post the kwargs to the flags API endpoint with automatic retries."""
res = post(
api_key,
host,
"/flags/?v=2",
gzip,
timeout,
session=_get_flags_session(),
**kwargs,
)
return _process_response(
res, success_message="Feature flags evaluated successfully"
)
def remote_config(
personal_api_key: str,
project_api_key: str,
host: Optional[str] = None,
key: str = "",
timeout: int = 15,
) -> Any:
"""Get remote config flag value from remote_config API endpoint"""
response = get(
personal_api_key,
f"/api/projects/@current/feature_flags/{key}/remote_config?token={project_api_key}",
host,
timeout,
)
return response.data
def batch_post(
api_key: str,
host: Optional[str] = None,
gzip: bool = False,
timeout: int = 15,
**kwargs,
) -> requests.Response:
"""Post the `kwargs` to the batch API endpoint for events"""
res = post(api_key, host, "/batch/", gzip, timeout, **kwargs)
return _process_response(
res, success_message="data uploaded successfully", return_json=False
)
def get(
api_key: str,
url: str,
host: Optional[str] = None,
timeout: Optional[int] = None,
etag: Optional[str] = None,
) -> GetResponse:
"""
Make a GET request with optional ETag support.
If an etag is provided, sends If-None-Match header. Returns GetResponse with:
- not_modified=True and data=None if server returns 304
- not_modified=False and data=response if server returns 200
"""
log = logging.getLogger("hanzo_insights")
full_url = remove_trailing_slash(host or DEFAULT_HOST) + url
headers = {"Authorization": "Bearer %s" % api_key, "User-Agent": USER_AGENT}
if etag:
headers["If-None-Match"] = etag
res = _get_session().get(full_url, headers=headers, timeout=timeout)
masked_url = _mask_tokens_in_url(full_url)
# Handle 304 Not Modified
if res.status_code == 304:
log.debug(f"GET {masked_url} returned 304 Not Modified")
response_etag = res.headers.get("ETag")
return GetResponse(data=None, etag=response_etag or etag, not_modified=True)
# Handle normal response
data = _process_response(
res, success_message=f"GET {masked_url} completed successfully"
)
response_etag = res.headers.get("ETag")
return GetResponse(data=data, etag=response_etag, not_modified=False)
class APIError(Exception):
def __init__(
self, status: Union[int, str], message: str, retry_after: Optional[float] = None
):
self.message = message
self.status = status
self.retry_after = retry_after
def __str__(self):
msg = "[Insights] {0} ({1})"
return msg.format(self.message, self.status)
class QuotaLimitError(APIError):
pass
# Re-export requests exceptions for use in client.py
# This keeps all requests library imports centralized in this module
RequestsTimeout = requests.exceptions.Timeout
RequestsConnectionError = requests.exceptions.ConnectionError
class DatetimeSerializer(json.JSONEncoder):
def default(self, obj: Any):
if isinstance(obj, (date, datetime)):
return obj.isoformat()
return json.JSONEncoder.default(self, obj)
@@ -6,7 +6,7 @@ import unittest
def all_names():
for _, modname, _ in pkgutil.iter_modules(__path__):
yield "posthog.test." + modname
yield "hanzo_insights.test." + modname
def all():
@@ -1,11 +1,14 @@
import json
from unittest.mock import patch
import pytest
from hanzo_insights import identify_context, new_context
try:
from anthropic.types import Message, Usage
from posthog.ai.anthropic import Anthropic, AsyncAnthropic
from hanzo_insights.ai.anthropic import Anthropic, AsyncAnthropic
ANTHROPIC_AVAILABLE = True
except ImportError:
@@ -102,7 +105,7 @@ class MockDelta:
@pytest.fixture
def mock_client():
with patch("posthog.client.Client") as mock_client:
with patch("hanzo_insights.client.Client") as mock_client:
mock_client.privacy_mode = False
yield mock_client
@@ -276,12 +279,12 @@ 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)
client = Anthropic(api_key="test-key", insights_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"},
insights_distinct_id="test-id",
insights_properties={"foo": "bar"},
)
assert response == mock_anthropic_response
@@ -305,19 +308,46 @@ def test_basic_completion(mock_client, mock_anthropic_response):
assert props["$ai_output_tokens"] == 10
assert props["$ai_http_status"] == 200
assert props["foo"] == "bar"
assert props["$ai_tokens_source"] == "sdk"
assert isinstance(props["$ai_latency"], float)
# Verify raw usage metadata is passed for backend processing
assert "$ai_usage" in props
assert props["$ai_usage"] is not None
# Verify it's JSON-serializable
json.dumps(props["$ai_usage"])
# Verify it has expected structure
assert isinstance(props["$ai_usage"], dict)
assert "input_tokens" in props["$ai_usage"]
assert "output_tokens" in props["$ai_usage"]
def test_tokens_source_passthrough(mock_client, mock_anthropic_response):
with patch(
"anthropic.resources.Messages.create", return_value=mock_anthropic_response
):
client = Anthropic(api_key="test-key", insights_client=mock_client)
client.messages.create(
model="claude-3-opus-20240229",
messages=[{"role": "user", "content": "Hello"}],
insights_distinct_id="test-id",
insights_properties={"$ai_input_tokens": 99999},
)
props = mock_client.capture.call_args[1]["properties"]
assert props["$ai_tokens_source"] == "passthrough"
assert props["$ai_input_tokens"] == 99999
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)
client = Anthropic(api_key="test-key", insights_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"},
insights_distinct_id="test-id",
insights_groups={"company": "test_company"},
)
assert response == mock_anthropic_response
@@ -331,12 +361,12 @@ 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)
client = Anthropic(api_key="test-key", insights_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,
insights_distinct_id="test-id",
insights_privacy_mode=True,
)
assert response == mock_anthropic_response
@@ -353,12 +383,12 @@ def test_privacy_mode_global(mock_client, mock_anthropic_response):
"anthropic.resources.Messages.create", return_value=mock_anthropic_response
):
mock_client.privacy_mode = True
client = Anthropic(api_key="test-key", posthog_client=mock_client)
client = Anthropic(api_key="test-key", insights_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,
insights_distinct_id="test-id",
insights_privacy_mode=False,
)
assert response == mock_anthropic_response
@@ -377,14 +407,14 @@ def test_basic_integration(mock_client):
"anthropic.resources.Messages.create",
return_value=create_mock_response(),
):
client = Anthropic(posthog_client=mock_client)
client = Anthropic(insights_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"},
insights_distinct_id="test-id",
insights_properties={"foo": "bar"},
system="You must always answer with 'Bar'.",
)
@@ -422,7 +452,7 @@ async def test_basic_async_integration(mock_client):
"anthropic.resources.messages.AsyncMessages.create",
side_effect=mock_async_create,
):
client = AsyncAnthropic(posthog_client=mock_client)
client = AsyncAnthropic(insights_client=mock_client)
await client.messages.create(
model="claude-3-opus-20240229",
messages=[
@@ -430,8 +460,8 @@ async def test_basic_async_integration(mock_client):
],
max_tokens=1,
temperature=0,
posthog_distinct_id="test-id",
posthog_properties={"foo": "bar"},
insights_distinct_id="test-id",
insights_properties={"foo": "bar"},
)
assert mock_client.capture.call_count == 1
@@ -483,7 +513,7 @@ async def test_async_streaming_system_prompt(mock_client):
"anthropic.resources.messages.AsyncMessages.create",
side_effect=async_create_wrapper,
):
client = AsyncAnthropic(posthog_client=mock_client)
client = AsyncAnthropic(insights_client=mock_client)
response = await client.messages.create(
model="claude-3-opus-20240229",
system="You must always answer with 'Bar'.",
@@ -511,7 +541,7 @@ def test_error(mock_client, mock_anthropic_response):
with patch(
"anthropic.resources.Messages.create", side_effect=Exception("Test error")
):
client = Anthropic(api_key="test-key", posthog_client=mock_client)
client = Anthropic(api_key="test-key", insights_client=mock_client)
with pytest.raises(Exception):
client.messages.create(
model="claude-3-opus-20240229",
@@ -531,12 +561,12 @@ def test_cached_tokens(mock_client, mock_anthropic_response_with_cached_tokens):
"anthropic.resources.Messages.create",
return_value=mock_anthropic_response_with_cached_tokens,
):
client = Anthropic(api_key="test-key", posthog_client=mock_client)
client = Anthropic(api_key="test-key", insights_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"},
insights_distinct_id="test-id",
insights_properties={"foo": "bar"},
)
assert response == mock_anthropic_response_with_cached_tokens
@@ -570,7 +600,7 @@ def test_tool_definition(mock_client, mock_anthropic_response):
"anthropic.resources.Messages.create",
return_value=mock_anthropic_response,
):
client = Anthropic(api_key="test-key", posthog_client=mock_client)
client = Anthropic(api_key="test-key", insights_client=mock_client)
tools = [
{
@@ -595,8 +625,8 @@ def test_tool_definition(mock_client, mock_anthropic_response):
temperature=0.7,
tools=tools,
messages=[{"role": "user", "content": "hey"}],
posthog_distinct_id="test-id",
posthog_properties={"foo": "bar"},
insights_distinct_id="test-id",
insights_properties={"foo": "bar"},
)
assert response == mock_anthropic_response
@@ -632,7 +662,7 @@ def test_tool_calls_in_output_choices(
"anthropic.resources.Messages.create",
return_value=mock_anthropic_response_with_tool_calls,
):
client = Anthropic(api_key="test-key", posthog_client=mock_client)
client = Anthropic(api_key="test-key", insights_client=mock_client)
response = client.messages.create(
model="claude-3-5-sonnet-20241022",
max_tokens=200,
@@ -650,7 +680,7 @@ def test_tool_calls_in_output_choices(
},
}
],
posthog_distinct_id="test-id",
insights_distinct_id="test-id",
)
assert response == mock_anthropic_response_with_tool_calls
@@ -694,7 +724,7 @@ def test_tool_calls_only_no_content(
"anthropic.resources.Messages.create",
return_value=mock_anthropic_response_tool_calls_only,
):
client = Anthropic(api_key="test-key", posthog_client=mock_client)
client = Anthropic(api_key="test-key", insights_client=mock_client)
response = client.messages.create(
model="claude-3-5-sonnet-20241022",
max_tokens=200,
@@ -713,7 +743,7 @@ def test_tool_calls_only_no_content(
},
}
],
posthog_distinct_id="test-id",
insights_distinct_id="test-id",
)
assert response == mock_anthropic_response_tool_calls_only
@@ -760,7 +790,7 @@ def test_async_tool_calls_in_output_choices(
"anthropic.resources.AsyncMessages.create",
side_effect=mock_async_create,
):
async_client = AsyncAnthropic(api_key="test-key", posthog_client=mock_client)
async_client = AsyncAnthropic(api_key="test-key", insights_client=mock_client)
async def run_test():
return await async_client.messages.create(
@@ -780,7 +810,7 @@ def test_async_tool_calls_in_output_choices(
},
}
],
posthog_distinct_id="test-id",
insights_distinct_id="test-id",
)
response = asyncio.run(run_test())
@@ -825,7 +855,7 @@ def test_streaming_with_tool_calls(mock_client, mock_anthropic_stream_with_tools
"anthropic.resources.Messages.create",
return_value=mock_anthropic_stream_with_tools,
):
client = Anthropic(api_key="test-key", posthog_client=mock_client)
client = Anthropic(api_key="test-key", insights_client=mock_client)
response = client.messages.create(
model="claude-3-5-sonnet-20241022",
system="You are a helpful weather assistant.",
@@ -847,7 +877,7 @@ def test_streaming_with_tool_calls(mock_client, mock_anthropic_stream_with_tools
}
],
stream=True,
posthog_distinct_id="test-id",
insights_distinct_id="test-id",
)
# Consume the stream - this triggers the finally block synchronously
@@ -917,6 +947,17 @@ def test_streaming_with_tool_calls(mock_client, mock_anthropic_stream_with_tools
assert props["$ai_output_tokens"] == 25
assert props["$ai_cache_read_input_tokens"] == 5
assert props["$ai_cache_creation_input_tokens"] == 0
assert props["$ai_tokens_source"] == "sdk"
# Verify raw usage is captured in streaming mode (merged from events)
assert "$ai_usage" in props
assert props["$ai_usage"] is not None
# Verify it's JSON-serializable
json.dumps(props["$ai_usage"])
# Verify it has expected structure (merged from message_start and message_delta)
assert isinstance(props["$ai_usage"], dict)
assert "input_tokens" in props["$ai_usage"]
assert "output_tokens" in props["$ai_usage"]
def test_async_streaming_with_tool_calls(mock_client, mock_anthropic_stream_with_tools):
@@ -936,7 +977,7 @@ def test_async_streaming_with_tool_calls(mock_client, mock_anthropic_stream_with
"anthropic.resources.AsyncMessages.create",
side_effect=mock_async_create,
):
async_client = AsyncAnthropic(api_key="test-key", posthog_client=mock_client)
async_client = AsyncAnthropic(api_key="test-key", insights_client=mock_client)
async def run_test():
response = await async_client.messages.create(
@@ -960,7 +1001,7 @@ def test_async_streaming_with_tool_calls(mock_client, mock_anthropic_stream_with
}
],
stream=True,
posthog_distinct_id="test-id",
insights_distinct_id="test-id",
)
# Consume the async stream
@@ -1060,11 +1101,11 @@ def test_web_search_count(mock_client):
mock_response = MockResponseWithWebSearch()
with patch("anthropic.resources.Messages.create", return_value=mock_response):
client = Anthropic(api_key="test-key", posthog_client=mock_client)
client = Anthropic(api_key="test-key", insights_client=mock_client)
response = client.messages.create(
model="claude-3-opus-20240229",
messages=[{"role": "user", "content": "Search for recent news"}],
posthog_distinct_id="test-id",
insights_distinct_id="test-id",
)
assert response == mock_response
@@ -1138,12 +1179,12 @@ def test_streaming_with_web_search(mock_client, mock_anthropic_stream_with_web_s
"anthropic.resources.Messages.create",
return_value=mock_anthropic_stream_with_web_search,
):
client = Anthropic(api_key="test-key", posthog_client=mock_client)
client = Anthropic(api_key="test-key", insights_client=mock_client)
response = client.messages.create(
model="claude-3-opus-20240229",
messages=[{"role": "user", "content": "Search for recent news"}],
stream=True,
posthog_distinct_id="test-id",
insights_distinct_id="test-id",
)
# Consume the stream - this triggers the finally block synchronously
@@ -1193,13 +1234,13 @@ def test_async_with_web_search(mock_client):
"anthropic.resources.AsyncMessages.create",
side_effect=mock_async_create,
):
async_client = AsyncAnthropic(api_key="test-key", posthog_client=mock_client)
async_client = AsyncAnthropic(api_key="test-key", insights_client=mock_client)
async def run_test():
response = await async_client.messages.create(
model="claude-3-opus-20240229",
messages=[{"role": "user", "content": "Search for recent news"}],
posthog_distinct_id="test-id",
insights_distinct_id="test-id",
)
return response
@@ -1237,14 +1278,14 @@ def test_async_streaming_with_web_search(
"anthropic.resources.AsyncMessages.create",
side_effect=mock_async_create,
):
async_client = AsyncAnthropic(api_key="test-key", posthog_client=mock_client)
async_client = AsyncAnthropic(api_key="test-key", insights_client=mock_client)
async def run_test():
response = await async_client.messages.create(
model="claude-3-opus-20240229",
messages=[{"role": "user", "content": "Search for recent news"}],
stream=True,
posthog_distinct_id="test-id",
insights_distinct_id="test-id",
)
# Consume the async stream
@@ -1263,3 +1304,99 @@ def test_async_streaming_with_web_search(
assert props["$ai_web_search_count"] == 2
assert props["$ai_input_tokens"] == 50
assert props["$ai_output_tokens"] == 25
# =======================
# Distinct ID Context Tests
# =======================
def test_no_distinct_id_uses_trace_id_and_personless(
mock_client, mock_anthropic_response
):
"""When no distinct_id is provided and no outer context, trace_id is used and event is personless."""
with patch(
"anthropic.resources.Messages.create", return_value=mock_anthropic_response
):
client = Anthropic(api_key="test-key", insights_client=mock_client)
client.messages.create(
model="claude-3-opus-20240229",
messages=[{"role": "user", "content": "Hello"}],
insights_trace_id="trace-123",
)
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
assert call_args["distinct_id"] == "trace-123"
assert props["$process_person_profile"] is False
def test_explicit_distinct_id_creates_person_profile(
mock_client, mock_anthropic_response
):
"""When insights_distinct_id is explicitly passed, it is used and event is not personless."""
with patch(
"anthropic.resources.Messages.create", return_value=mock_anthropic_response
):
client = Anthropic(api_key="test-key", insights_client=mock_client)
client.messages.create(
model="claude-3-opus-20240229",
messages=[{"role": "user", "content": "Hello"}],
insights_distinct_id="user-123",
insights_trace_id="trace-123",
)
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
assert call_args["distinct_id"] == "user-123"
assert (
"$process_person_profile" not in props
or props["$process_person_profile"] is not False
)
def test_outer_context_distinct_id_is_used(mock_client, mock_anthropic_response):
"""When an outer context has a distinct_id, it should be used instead of trace_id."""
with patch(
"anthropic.resources.Messages.create", return_value=mock_anthropic_response
):
client = Anthropic(api_key="test-key", insights_client=mock_client)
with new_context():
identify_context("outer-user-456")
client.messages.create(
model="claude-3-opus-20240229",
messages=[{"role": "user", "content": "Hello"}],
insights_trace_id="trace-123",
)
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
assert call_args["distinct_id"] == "outer-user-456"
assert (
"$process_person_profile" not in props
or props["$process_person_profile"] is not False
)
def test_explicit_distinct_id_overrides_outer_context(
mock_client, mock_anthropic_response
):
"""When both outer context and explicit insights_distinct_id are set, explicit wins."""
with patch(
"anthropic.resources.Messages.create", return_value=mock_anthropic_response
):
client = Anthropic(api_key="test-key", insights_client=mock_client)
with new_context():
identify_context("outer-user-456")
client.messages.create(
model="claude-3-opus-20240229",
messages=[{"role": "user", "content": "Hello"}],
insights_distinct_id="explicit-user-789",
insights_trace_id="trace-123",
)
call_args = mock_client.capture.call_args[1]
assert call_args["distinct_id"] == "explicit-user-789"
@@ -1,3 +1,4 @@
import json
from unittest.mock import MagicMock, patch
import pytest
@@ -5,7 +6,7 @@ import pytest
try:
from google import genai as google_genai
from posthog.ai.gemini import Client
from hanzo_insights.ai.gemini import Client
GEMINI_AVAILABLE = True
except ImportError:
@@ -18,7 +19,7 @@ pytestmark = pytest.mark.skipif(
@pytest.fixture
def mock_client():
with patch("posthog.client.Client") as mock_client:
with patch("hanzo_insights.client.Client") as mock_client:
mock_client.privacy_mode = False
yield mock_client
@@ -34,6 +35,13 @@ def mock_gemini_response():
# Ensure cache and reasoning tokens are not present (not MagicMock)
mock_usage.cached_content_token_count = 0
mock_usage.thoughts_token_count = 0
# Make model_dump() return a proper dict for serialization
mock_usage.model_dump.return_value = {
"prompt_token_count": 20,
"candidates_token_count": 10,
"cached_content_token_count": 0,
"thoughts_token_count": 0,
}
mock_response.usage_metadata = mock_usage
mock_candidate = MagicMock()
@@ -69,6 +77,13 @@ def mock_gemini_response_with_function_calls():
mock_usage.candidates_token_count = 15
mock_usage.cached_content_token_count = 0
mock_usage.thoughts_token_count = 0
# Make model_dump() return a proper dict for serialization
mock_usage.model_dump.return_value = {
"prompt_token_count": 25,
"candidates_token_count": 15,
"cached_content_token_count": 0,
"thoughts_token_count": 0,
}
mock_response.usage_metadata = mock_usage
# Mock function call
@@ -117,6 +132,13 @@ def mock_gemini_response_function_calls_only():
mock_usage.candidates_token_count = 12
mock_usage.cached_content_token_count = 0
mock_usage.thoughts_token_count = 0
# Make model_dump() return a proper dict for serialization
mock_usage.model_dump.return_value = {
"prompt_token_count": 30,
"candidates_token_count": 12,
"cached_content_token_count": 0,
"thoughts_token_count": 0,
}
mock_response.usage_metadata = mock_usage
# Mock function call
@@ -150,13 +172,13 @@ def test_new_client_basic_generation(
"""Test the new Client/Models API structure"""
mock_google_genai_client.models.generate_content.return_value = mock_gemini_response
client = Client(api_key="test-key", posthog_client=mock_client)
client = Client(api_key="test-key", insights_client=mock_client)
response = client.models.generate_content(
model="gemini-2.0-flash",
contents=["Tell me a fun fact about hedgehogs"],
posthog_distinct_id="test-id",
posthog_properties={"foo": "bar"},
insights_distinct_id="test-id",
insights_properties={"foo": "bar"},
)
assert response == mock_gemini_response
@@ -174,6 +196,15 @@ def test_new_client_basic_generation(
assert props["foo"] == "bar"
assert "$ai_trace_id" in props
assert props["$ai_latency"] > 0
# Verify raw usage metadata is passed for backend processing
assert "$ai_usage" in props
assert props["$ai_usage"] is not None
# Verify it's JSON-serializable
json.dumps(props["$ai_usage"])
# Verify it has expected structure
assert isinstance(props["$ai_usage"], dict)
assert "prompt_token_count" in props["$ai_usage"]
assert "candidates_token_count" in props["$ai_usage"]
def test_new_client_streaming_with_generate_content_stream(
@@ -208,13 +239,13 @@ def test_new_client_streaming_with_generate_content_stream(
mock_streaming_response()
)
client = Client(api_key="test-key", posthog_client=mock_client)
client = Client(api_key="test-key", insights_client=mock_client)
response = client.models.generate_content_stream(
model="gemini-2.0-flash",
contents=["Write a short story"],
posthog_distinct_id="test-id",
posthog_properties={"feature": "streaming"},
insights_distinct_id="test-id",
insights_properties={"feature": "streaming"},
)
chunks = list(response)
@@ -267,7 +298,7 @@ def test_new_client_streaming_with_tools(mock_client, mock_google_genai_client):
mock_streaming_response()
)
client = Client(api_key="test-key", posthog_client=mock_client)
client = Client(api_key="test-key", insights_client=mock_client)
# Create mock tools configuration
mock_tool = MagicMock()
@@ -295,8 +326,8 @@ def test_new_client_streaming_with_tools(mock_client, mock_google_genai_client):
model="gemini-2.0-flash",
contents=["What's the weather in SF?"],
config=mock_config,
posthog_distinct_id="test-id",
posthog_properties={"feature": "streaming_with_tools"},
insights_distinct_id="test-id",
insights_properties={"feature": "streaming_with_tools"},
)
chunks = list(response)
@@ -326,13 +357,13 @@ def test_new_client_groups(mock_client, mock_google_genai_client, mock_gemini_re
"""Test groups functionality with new Client API"""
mock_google_genai_client.models.generate_content.return_value = mock_gemini_response
client = Client(api_key="test-key", posthog_client=mock_client)
client = Client(api_key="test-key", insights_client=mock_client)
client.models.generate_content(
model="gemini-2.0-flash",
contents=["Hello"],
posthog_distinct_id="test-id",
posthog_groups={"company": "company_123"},
insights_distinct_id="test-id",
insights_groups={"company": "company_123"},
)
call_args = mock_client.capture.call_args[1]
@@ -345,13 +376,13 @@ def test_new_client_privacy_mode_local(
"""Test local privacy mode with new Client API"""
mock_google_genai_client.models.generate_content.return_value = mock_gemini_response
client = Client(api_key="test-key", posthog_client=mock_client)
client = Client(api_key="test-key", insights_client=mock_client)
client.models.generate_content(
model="gemini-2.0-flash",
contents=["Hello"],
posthog_distinct_id="test-id",
posthog_privacy_mode=True,
insights_distinct_id="test-id",
insights_privacy_mode=True,
)
call_args = mock_client.capture.call_args[1]
@@ -368,12 +399,12 @@ def test_new_client_privacy_mode_global(
mock_google_genai_client.models.generate_content.return_value = mock_gemini_response
client = Client(api_key="test-key", posthog_client=mock_client)
client = Client(api_key="test-key", insights_client=mock_client)
client.models.generate_content(
model="gemini-2.0-flash",
contents=["Hello"],
posthog_distinct_id="test-id",
insights_distinct_id="test-id",
)
call_args = mock_client.capture.call_args[1]
@@ -388,11 +419,11 @@ def test_new_client_different_input_formats(
"""Test different input formats with new Client API"""
mock_google_genai_client.models.generate_content.return_value = mock_gemini_response
client = Client(api_key="test-key", posthog_client=mock_client)
client = Client(api_key="test-key", insights_client=mock_client)
# Test string input
client.models.generate_content(
model="gemini-2.0-flash", contents="Hello", posthog_distinct_id="test-id"
model="gemini-2.0-flash", contents="Hello", insights_distinct_id="test-id"
)
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
@@ -403,27 +434,37 @@ def test_new_client_different_input_formats(
client.models.generate_content(
model="gemini-2.0-flash",
contents=[{"role": "user", "parts": [{"text": "hey"}]}],
posthog_distinct_id="test-id",
insights_distinct_id="test-id",
)
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
assert props["$ai_input"] == [{"role": "user", "content": "hey"}]
assert props["$ai_input"] == [
{"role": "user", "content": [{"type": "text", "text": "hey"}]}
]
# Test multiple parts in the parts array
mock_client.reset_mock()
client.models.generate_content(
model="gemini-2.0-flash",
contents=[{"role": "user", "parts": [{"text": "Hello "}, {"text": "world"}]}],
posthog_distinct_id="test-id",
insights_distinct_id="test-id",
)
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
assert props["$ai_input"] == [{"role": "user", "content": "Hello world"}]
assert props["$ai_input"] == [
{
"role": "user",
"content": [
{"type": "text", "text": "Hello "},
{"type": "text", "text": "world"},
],
}
]
# Test list input with string
mock_client.capture.reset_mock()
client.models.generate_content(
model="gemini-2.0-flash", contents=["List item"], posthog_distinct_id="test-id"
model="gemini-2.0-flash", contents=["List item"], insights_distinct_id="test-id"
)
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
@@ -436,12 +477,12 @@ def test_new_client_model_parameters(
"""Test model parameters with new Client API"""
mock_google_genai_client.models.generate_content.return_value = mock_gemini_response
client = Client(api_key="test-key", posthog_client=mock_client)
client = Client(api_key="test-key", insights_client=mock_client)
client.models.generate_content(
model="gemini-2.0-flash",
contents=["Hello"],
posthog_distinct_id="test-id",
insights_distinct_id="test-id",
temperature=0.7,
max_tokens=100,
)
@@ -455,16 +496,16 @@ def test_new_client_model_parameters(
def test_new_client_default_settings(
mock_client, mock_google_genai_client, mock_gemini_response
):
"""Test client with default PostHog settings"""
"""Test client with default Insights settings"""
mock_google_genai_client.models.generate_content.return_value = mock_gemini_response
client = Client(
api_key="test-key",
posthog_client=mock_client,
posthog_distinct_id="default_user",
posthog_properties={"team": "ai"},
posthog_privacy_mode=False,
posthog_groups={"company": "acme_corp"},
insights_client=mock_client,
insights_distinct_id="default_user",
insights_properties={"team": "ai"},
insights_privacy_mode=False,
insights_groups={"company": "acme_corp"},
)
# Call without overriding defaults
@@ -486,21 +527,21 @@ def test_new_client_override_defaults(
client = Client(
api_key="test-key",
posthog_client=mock_client,
posthog_distinct_id="default_user",
posthog_properties={"team": "ai"},
posthog_privacy_mode=False,
posthog_groups={"company": "acme_corp"},
insights_client=mock_client,
insights_distinct_id="default_user",
insights_properties={"team": "ai"},
insights_privacy_mode=False,
insights_groups={"company": "acme_corp"},
)
# Override defaults in call
client.models.generate_content(
model="gemini-2.0-flash",
contents=["Hello"],
posthog_distinct_id="specific_user",
posthog_properties={"feature": "chat", "urgent": True},
posthog_privacy_mode=True,
posthog_groups={"organization": "special_org"},
insights_distinct_id="specific_user",
insights_properties={"feature": "chat", "urgent": True},
insights_privacy_mode=True,
insights_groups={"organization": "special_org"},
)
call_args = mock_client.capture.call_args[1]
@@ -536,7 +577,7 @@ def test_vertex_ai_parameters_passed_through(
location="us-central1",
debug_config=mock_debug_config,
http_options=mock_http_options,
posthog_client=mock_client,
insights_client=mock_client,
)
# Verify genai.Client was called with correct parameters
@@ -556,7 +597,7 @@ def test_api_key_mode(mock_client, mock_google_genai_client):
# Create client with just API key (traditional mode)
Client(
api_key="test-api-key",
posthog_client=mock_client,
insights_client=mock_client,
)
# Verify genai.Client was called with only api_key
@@ -577,7 +618,7 @@ def test_vertex_ai_mode_with_optional_api_key(
api_key="test-api-key",
credentials=mock_credentials,
project="test-project",
posthog_client=mock_client,
insights_client=mock_client,
)
# Verify genai.Client was called with both Vertex AI params and API key
@@ -593,7 +634,7 @@ def test_tool_use_response(mock_client, mock_google_genai_client, mock_gemini_re
"""Test that tools defined in config are captured in $ai_tools property"""
mock_google_genai_client.models.generate_content.return_value = mock_gemini_response
client = Client(api_key="test-key", posthog_client=mock_client)
client = Client(api_key="test-key", insights_client=mock_client)
# Create mock tools configuration
mock_tool = MagicMock()
@@ -623,8 +664,8 @@ def test_tool_use_response(mock_client, mock_google_genai_client, mock_gemini_re
model="gemini-2.5-flash",
contents=["hey"],
config=mock_config,
posthog_distinct_id="test-id",
posthog_properties={"foo": "bar"},
insights_distinct_id="test-id",
insights_properties={"foo": "bar"},
)
assert response == mock_gemini_response
@@ -661,12 +702,12 @@ def test_function_calls_in_output_choices(
mock_gemini_response_with_function_calls
)
client = Client(api_key="test-key", posthog_client=mock_client)
client = Client(api_key="test-key", insights_client=mock_client)
response = client.models.generate_content(
model="gemini-2.5-flash",
contents=["What's the weather in San Francisco?"],
posthog_distinct_id="test-id",
insights_distinct_id="test-id",
)
assert response == mock_gemini_response_with_function_calls
@@ -710,12 +751,12 @@ def test_function_calls_only_no_content(
mock_gemini_response_function_calls_only
)
client = Client(api_key="test-key", posthog_client=mock_client)
client = Client(api_key="test-key", insights_client=mock_client)
response = client.models.generate_content(
model="gemini-2.5-flash",
contents=["Get weather for New York"],
posthog_distinct_id="test-id",
insights_distinct_id="test-id",
)
assert response == mock_gemini_response_function_calls_only
@@ -769,12 +810,12 @@ def test_cache_and_reasoning_tokens(mock_client, mock_google_genai_client):
mock_google_genai_client.models.generate_content.return_value = mock_response
client = Client(api_key="test-key", posthog_client=mock_client)
client = Client(api_key="test-key", insights_client=mock_client)
response = client.models.generate_content(
model="gemini-2.5-pro",
contents="Test with cache",
posthog_distinct_id="test-id",
insights_distinct_id="test-id",
)
assert response == mock_response
@@ -800,6 +841,13 @@ def test_streaming_cache_and_reasoning_tokens(mock_client, mock_google_genai_cli
chunk1_usage.candidates_token_count = 5
chunk1_usage.cached_content_token_count = 30 # Cache tokens
chunk1_usage.thoughts_token_count = 0
# Make model_dump() return a proper dict for serialization
chunk1_usage.model_dump.return_value = {
"prompt_token_count": 100,
"candidates_token_count": 5,
"cached_content_token_count": 30,
"thoughts_token_count": 0,
}
chunk1.usage_metadata = chunk1_usage
chunk2 = MagicMock()
@@ -809,24 +857,31 @@ def test_streaming_cache_and_reasoning_tokens(mock_client, mock_google_genai_cli
chunk2_usage.candidates_token_count = 10
chunk2_usage.cached_content_token_count = 30 # Same cache tokens
chunk2_usage.thoughts_token_count = 5 # Reasoning tokens
# Make model_dump() return a proper dict for serialization
chunk2_usage.model_dump.return_value = {
"prompt_token_count": 100,
"candidates_token_count": 10,
"cached_content_token_count": 30,
"thoughts_token_count": 5,
}
chunk2.usage_metadata = chunk2_usage
mock_stream = iter([chunk1, chunk2])
mock_google_genai_client.models.generate_content_stream.return_value = mock_stream
client = Client(api_key="test-key", posthog_client=mock_client)
client = Client(api_key="test-key", insights_client=mock_client)
response = client.models.generate_content_stream(
model="gemini-2.5-pro",
contents="Test streaming with cache",
posthog_distinct_id="test-id",
insights_distinct_id="test-id",
)
# Consume the stream
result = list(response)
assert len(result) == 2
# Check PostHog capture was called
# Check Insights capture was called
assert mock_client.capture.call_count == 1
call_args = mock_client.capture.call_args[1]
@@ -838,6 +893,16 @@ def test_streaming_cache_and_reasoning_tokens(mock_client, mock_google_genai_cli
assert props["$ai_cache_read_input_tokens"] == 30
assert props["$ai_reasoning_tokens"] == 5
# Verify raw usage is captured in streaming mode (merged from chunks)
assert "$ai_usage" in props
assert props["$ai_usage"] is not None
# Verify it's JSON-serializable
json.dumps(props["$ai_usage"])
# Verify it has expected structure
assert isinstance(props["$ai_usage"], dict)
assert "prompt_token_count" in props["$ai_usage"]
assert "candidates_token_count" in props["$ai_usage"]
def test_web_search_grounding(mock_client, mock_google_genai_client):
"""Test web search detection via grounding_metadata."""
@@ -881,11 +946,11 @@ def test_web_search_grounding(mock_client, mock_google_genai_client):
# Mock the generate_content method
mock_google_genai_client.models.generate_content.return_value = mock_response
client = Client(api_key="test-key", posthog_client=mock_client)
client = Client(api_key="test-key", insights_client=mock_client)
response = client.models.generate_content(
model="gemini-2.5-flash",
contents="What's the latest news?",
posthog_distinct_id="test-id",
insights_distinct_id="test-id",
)
assert response == mock_response
@@ -950,12 +1015,12 @@ def test_streaming_with_web_search(mock_client, mock_google_genai_client):
mock_streaming_response()
)
client = Client(api_key="test-key", posthog_client=mock_client)
client = Client(api_key="test-key", insights_client=mock_client)
response = client.models.generate_content_stream(
model="gemini-2.5-flash",
contents="What's the latest news?",
posthog_distinct_id="test-id",
insights_distinct_id="test-id",
)
chunks = list(response)
@@ -1015,12 +1080,12 @@ def test_empty_grounding_metadata_no_web_search(mock_client, mock_google_genai_c
# Mock the generate_content method
mock_google_genai_client.models.generate_content.return_value = mock_response
client = Client(api_key="test-key", posthog_client=mock_client)
client = Client(api_key="test-key", insights_client=mock_client)
response = client.models.generate_content(
model="gemini-2.5-flash",
contents="Hello",
posthog_distinct_id="test-id",
insights_distinct_id="test-id",
)
assert response == mock_response
@@ -1078,12 +1143,12 @@ def test_empty_array_grounding_metadata_no_web_search(
# Mock the generate_content method
mock_google_genai_client.models.generate_content.return_value = mock_response
client = Client(api_key="test-key", posthog_client=mock_client)
client = Client(api_key="test-key", insights_client=mock_client)
response = client.models.generate_content(
model="gemini-2.5-flash",
contents="What can you do?",
posthog_distinct_id="test-id",
insights_distinct_id="test-id",
)
assert response == mock_response
@@ -0,0 +1,853 @@
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
try:
from google import genai as google_genai
from hanzo_insights.ai.gemini import AsyncClient
GEMINI_AVAILABLE = True
except ImportError:
GEMINI_AVAILABLE = False
pytestmark = [
pytest.mark.skipif(
not GEMINI_AVAILABLE, reason="Google Gemini package is not available"
),
pytest.mark.asyncio,
]
@pytest.fixture
def mock_client():
with patch("hanzo_insights.client.Client") as mock_client:
mock_client.privacy_mode = False
yield mock_client
@pytest.fixture
def mock_gemini_response():
mock_response = MagicMock()
mock_response.text = "Test response from Gemini"
mock_usage = MagicMock()
mock_usage.prompt_token_count = 20
mock_usage.candidates_token_count = 10
# Ensure cache and reasoning tokens are not present (not MagicMock)
mock_usage.cached_content_token_count = 0
mock_usage.thoughts_token_count = 0
mock_response.usage_metadata = mock_usage
mock_candidate = MagicMock()
mock_candidate.text = "Test response from Gemini"
mock_content = MagicMock()
mock_part = MagicMock()
mock_part.text = "Test response from Gemini"
mock_content.parts = [mock_part]
mock_candidate.content = mock_content
mock_response.candidates = [mock_candidate]
return mock_response
@pytest.fixture
def mock_google_genai_client():
"""Mock for the google-genai Client with async support"""
with patch.object(google_genai, "Client") as mock_client_class:
mock_client_instance = MagicMock()
mock_models = MagicMock()
mock_aio = MagicMock()
mock_aio_models = MagicMock()
mock_client_instance.models = mock_models
mock_client_instance.aio = mock_aio
mock_aio.models = mock_aio_models
mock_client_class.return_value = mock_client_instance
yield mock_client_instance
@pytest.fixture
def mock_gemini_response_with_function_calls():
mock_response = MagicMock()
# Mock usage metadata
mock_usage = MagicMock()
mock_usage.prompt_token_count = 25
mock_usage.candidates_token_count = 15
mock_usage.cached_content_token_count = 0
mock_usage.thoughts_token_count = 0
mock_response.usage_metadata = mock_usage
# Mock function call
mock_function_call = MagicMock()
mock_function_call.name = "get_current_weather"
mock_function_call.args = {"location": "San Francisco"}
# Mock text part 1
mock_text_part1 = MagicMock()
mock_text_part1.text = "I'll check the weather for you."
type(mock_text_part1).text = mock_text_part1.text
# Mock text part 2
mock_text_part2 = MagicMock()
mock_text_part2.text = " Let me look that up."
type(mock_text_part2).text = mock_text_part2.text
# Mock function call part
mock_function_part = MagicMock()
mock_function_part.function_call = mock_function_call
type(mock_function_part).function_call = mock_function_part.function_call
del mock_function_part.text
# Mock content with 2 text parts and 1 function call part
mock_content = MagicMock()
mock_content.parts = [mock_text_part1, mock_text_part2, mock_function_part]
# Mock candidate
mock_candidate = MagicMock()
mock_candidate.content = mock_content
mock_response.candidates = [mock_candidate]
return mock_response
async def test_async_client_basic_generation(
mock_client, mock_google_genai_client, mock_gemini_response
):
"""Test the async Client/AsyncModels API structure"""
mock_google_genai_client.aio.models.generate_content = AsyncMock(
return_value=mock_gemini_response
)
client = AsyncClient(api_key="test-key", insights_client=mock_client)
response = await client.models.generate_content(
model="gemini-2.0-flash",
contents=["Tell me a fun fact about hedgehogs"],
insights_distinct_id="test-id",
insights_properties={"foo": "bar"},
)
assert response == mock_gemini_response
assert mock_client.capture.call_count == 1
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
assert call_args["distinct_id"] == "test-id"
assert call_args["event"] == "$ai_generation"
assert props["$ai_provider"] == "gemini"
assert props["$ai_model"] == "gemini-2.0-flash"
assert props["$ai_input_tokens"] == 20
assert props["$ai_output_tokens"] == 10
assert props["foo"] == "bar"
assert "$ai_trace_id" in props
assert props["$ai_latency"] > 0
async def test_async_client_streaming_with_generate_content_stream(
mock_client, mock_google_genai_client
):
"""Test the async generate_content_stream method"""
async def mock_streaming_response():
mock_chunk1 = MagicMock()
mock_chunk1.text = "Hello "
mock_usage1 = MagicMock()
mock_usage1.prompt_token_count = 10
mock_usage1.candidates_token_count = 5
mock_usage1.cached_content_token_count = 0
mock_usage1.thoughts_token_count = 0
mock_chunk1.usage_metadata = mock_usage1
yield mock_chunk1
mock_chunk2 = MagicMock()
mock_chunk2.text = "world!"
mock_usage2 = MagicMock()
mock_usage2.prompt_token_count = 10
mock_usage2.candidates_token_count = 10
mock_usage2.cached_content_token_count = 0
mock_usage2.thoughts_token_count = 0
mock_chunk2.usage_metadata = mock_usage2
yield mock_chunk2
# Mock the async generate_content_stream method
mock_google_genai_client.aio.models.generate_content_stream = AsyncMock(
return_value=mock_streaming_response()
)
client = AsyncClient(api_key="test-key", insights_client=mock_client)
response = await client.models.generate_content_stream(
model="gemini-2.0-flash",
contents=["Write a short story"],
insights_distinct_id="test-id",
insights_properties={"feature": "streaming"},
)
chunks = []
async for chunk in response:
chunks.append(chunk)
assert len(chunks) == 2
assert chunks[0].text == "Hello "
assert chunks[1].text == "world!"
# Check that the streaming event was captured
assert mock_client.capture.call_count == 1
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
assert call_args["distinct_id"] == "test-id"
assert call_args["event"] == "$ai_generation"
assert props["$ai_provider"] == "gemini"
assert props["$ai_model"] == "gemini-2.0-flash"
assert props["$ai_input_tokens"] == 10
assert props["$ai_output_tokens"] == 10
assert props["feature"] == "streaming"
assert isinstance(props["$ai_latency"], float)
async def test_async_client_streaming_with_tools(mock_client, mock_google_genai_client):
"""Test that tools are captured in async streaming mode"""
async def mock_streaming_response():
mock_chunk1 = MagicMock()
mock_chunk1.text = "I'll check "
mock_usage1 = MagicMock()
mock_usage1.prompt_token_count = 15
mock_usage1.candidates_token_count = 5
mock_usage1.cached_content_token_count = 0
mock_usage1.thoughts_token_count = 0
mock_chunk1.usage_metadata = mock_usage1
yield mock_chunk1
mock_chunk2 = MagicMock()
mock_chunk2.text = "the weather"
mock_usage2 = MagicMock()
mock_usage2.prompt_token_count = 15
mock_usage2.candidates_token_count = 10
mock_usage2.cached_content_token_count = 0
mock_usage2.thoughts_token_count = 0
mock_chunk2.usage_metadata = mock_usage2
yield mock_chunk2
# Mock the async generate_content_stream method
mock_google_genai_client.aio.models.generate_content_stream = AsyncMock(
return_value=mock_streaming_response()
)
client = AsyncClient(api_key="test-key", insights_client=mock_client)
# Create mock tools configuration
mock_tool = MagicMock()
mock_tool.function_declarations = [
MagicMock(
name="get_current_weather",
description="Gets the current weather for a given location.",
parameters=MagicMock(
type="OBJECT",
properties={
"location": MagicMock(
type="STRING",
description="The city and state, e.g. San Francisco, CA",
)
},
required=["location"],
),
)
]
mock_config = MagicMock()
mock_config.tools = [mock_tool]
response = await client.models.generate_content_stream(
model="gemini-2.0-flash",
contents=["What's the weather in SF?"],
config=mock_config,
insights_distinct_id="test-id",
insights_properties={"feature": "streaming_with_tools"},
)
chunks = []
async for chunk in response:
chunks.append(chunk)
assert len(chunks) == 2
assert chunks[0].text == "I'll check "
assert chunks[1].text == "the weather"
# Check that the streaming event was captured with tools
assert mock_client.capture.call_count == 1
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
assert call_args["distinct_id"] == "test-id"
assert call_args["event"] == "$ai_generation"
assert props["$ai_provider"] == "gemini"
assert props["$ai_model"] == "gemini-2.0-flash"
assert props["$ai_input_tokens"] == 15
assert props["$ai_output_tokens"] == 10
assert props["feature"] == "streaming_with_tools"
assert isinstance(props["$ai_latency"], float)
# Verify that tools are captured in the $ai_tools property in streaming mode
assert props["$ai_tools"] == [mock_tool]
async def test_async_client_groups(
mock_client, mock_google_genai_client, mock_gemini_response
):
"""Test groups functionality with async Client API"""
mock_google_genai_client.aio.models.generate_content = AsyncMock(
return_value=mock_gemini_response
)
client = AsyncClient(api_key="test-key", insights_client=mock_client)
await client.models.generate_content(
model="gemini-2.0-flash",
contents=["Hello"],
insights_distinct_id="test-id",
insights_groups={"company": "company_123"},
)
call_args = mock_client.capture.call_args[1]
assert call_args["groups"] == {"company": "company_123"}
async def test_async_client_privacy_mode_local(
mock_client, mock_google_genai_client, mock_gemini_response
):
"""Test local privacy mode with async Client API"""
mock_google_genai_client.aio.models.generate_content = AsyncMock(
return_value=mock_gemini_response
)
client = AsyncClient(api_key="test-key", insights_client=mock_client)
await client.models.generate_content(
model="gemini-2.0-flash",
contents=["Hello"],
insights_distinct_id="test-id",
insights_privacy_mode=True,
)
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
assert props["$ai_input"] is None
assert props["$ai_output_choices"] is None
async def test_async_client_privacy_mode_global(
mock_client, mock_google_genai_client, mock_gemini_response
):
"""Test global privacy mode with async Client API"""
mock_client.privacy_mode = True
mock_google_genai_client.aio.models.generate_content = AsyncMock(
return_value=mock_gemini_response
)
client = AsyncClient(api_key="test-key", insights_client=mock_client)
await client.models.generate_content(
model="gemini-2.0-flash",
contents=["Hello"],
insights_distinct_id="test-id",
)
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
assert props["$ai_input"] is None
assert props["$ai_output_choices"] is None
async def test_async_client_different_input_formats(
mock_client, mock_google_genai_client, mock_gemini_response
):
"""Test different input formats with async Client API"""
mock_google_genai_client.aio.models.generate_content = AsyncMock(
return_value=mock_gemini_response
)
client = AsyncClient(api_key="test-key", insights_client=mock_client)
# Test string input
await client.models.generate_content(
model="gemini-2.0-flash", contents="Hello", insights_distinct_id="test-id"
)
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
assert props["$ai_input"] == [{"role": "user", "content": "Hello"}]
# Test Gemini-specific format with parts array
mock_client.reset_mock()
await client.models.generate_content(
model="gemini-2.0-flash",
contents=[{"role": "user", "parts": [{"text": "hey"}]}],
insights_distinct_id="test-id",
)
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
assert props["$ai_input"] == [
{"role": "user", "content": [{"type": "text", "text": "hey"}]}
]
# Test multiple parts in the parts array
mock_client.reset_mock()
await client.models.generate_content(
model="gemini-2.0-flash",
contents=[{"role": "user", "parts": [{"text": "Hello "}, {"text": "world"}]}],
insights_distinct_id="test-id",
)
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
assert props["$ai_input"] == [
{
"role": "user",
"content": [
{"type": "text", "text": "Hello "},
{"type": "text", "text": "world"},
],
}
]
# Test list input with string
mock_client.capture.reset_mock()
await client.models.generate_content(
model="gemini-2.0-flash", contents=["List item"], insights_distinct_id="test-id"
)
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
assert props["$ai_input"] == [{"role": "user", "content": "List item"}]
async def test_async_client_model_parameters(
mock_client, mock_google_genai_client, mock_gemini_response
):
"""Test model parameters with async Client API"""
mock_google_genai_client.aio.models.generate_content = AsyncMock(
return_value=mock_gemini_response
)
client = AsyncClient(api_key="test-key", insights_client=mock_client)
await client.models.generate_content(
model="gemini-2.0-flash",
contents=["Hello"],
insights_distinct_id="test-id",
temperature=0.7,
max_tokens=100,
)
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
assert props["$ai_model_parameters"]["temperature"] == 0.7
assert props["$ai_model_parameters"]["max_tokens"] == 100
async def test_async_client_default_settings(
mock_client, mock_google_genai_client, mock_gemini_response
):
"""Test async client with default Insights settings"""
mock_google_genai_client.aio.models.generate_content = AsyncMock(
return_value=mock_gemini_response
)
client = AsyncClient(
api_key="test-key",
insights_client=mock_client,
insights_distinct_id="default_user",
insights_properties={"team": "ai"},
insights_privacy_mode=False,
insights_groups={"company": "acme_corp"},
)
# Call without overriding defaults
await client.models.generate_content(model="gemini-2.0-flash", contents=["Hello"])
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
assert call_args["distinct_id"] == "default_user"
assert call_args["groups"] == {"company": "acme_corp"}
assert props["team"] == "ai"
async def test_async_client_override_defaults(
mock_client, mock_google_genai_client, mock_gemini_response
):
"""Test overriding async client defaults per call"""
mock_google_genai_client.aio.models.generate_content = AsyncMock(
return_value=mock_gemini_response
)
client = AsyncClient(
api_key="test-key",
insights_client=mock_client,
insights_distinct_id="default_user",
insights_properties={"team": "ai"},
insights_privacy_mode=False,
insights_groups={"company": "acme_corp"},
)
# Override defaults in call
await client.models.generate_content(
model="gemini-2.0-flash",
contents=["Hello"],
insights_distinct_id="specific_user",
insights_properties={"feature": "chat", "urgent": True},
insights_privacy_mode=True,
insights_groups={"organization": "special_org"},
)
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
# Check overrides
assert call_args["distinct_id"] == "specific_user"
assert call_args["groups"] == {"organization": "special_org"}
assert props["$ai_input"] is None # privacy mode was overridden
# Check merged properties (defaults + call-specific)
assert props["team"] == "ai" # from defaults
assert props["feature"] == "chat" # from call
assert props["urgent"] is True # from call
async def test_async_vertex_ai_parameters_passed_through(
mock_client, mock_google_genai_client, mock_gemini_response
):
"""Test that Vertex AI parameters are properly passed to genai.Client"""
mock_google_genai_client.aio.models.generate_content = AsyncMock(
return_value=mock_gemini_response
)
# Mock credentials object
mock_credentials = MagicMock()
mock_debug_config = MagicMock()
mock_http_options = MagicMock()
# Create client with Vertex AI parameters
AsyncClient(
vertexai=True,
credentials=mock_credentials,
project="test-project",
location="us-central1",
debug_config=mock_debug_config,
http_options=mock_http_options,
insights_client=mock_client,
)
# Verify genai.Client was called with correct parameters
google_genai.Client.assert_called_once_with(
vertexai=True,
credentials=mock_credentials,
project="test-project",
location="us-central1",
debug_config=mock_debug_config,
http_options=mock_http_options,
)
async def test_async_api_key_mode(mock_client, mock_google_genai_client):
"""Test API key authentication mode with async client"""
# Create async client with just API key (traditional mode)
AsyncClient(
api_key="test-api-key",
insights_client=mock_client,
)
# Verify genai.Client was called with only api_key
google_genai.Client.assert_called_once_with(api_key="test-api-key")
async def test_async_function_calls_in_output_choices(
mock_client, mock_google_genai_client, mock_gemini_response_with_function_calls
):
"""Test that function calls are properly included in $ai_output_choices with async"""
mock_google_genai_client.aio.models.generate_content = AsyncMock(
return_value=mock_gemini_response_with_function_calls
)
client = AsyncClient(api_key="test-key", insights_client=mock_client)
response = await client.models.generate_content(
model="gemini-2.5-flash",
contents=["What's the weather in San Francisco?"],
insights_distinct_id="test-id",
)
assert response == mock_gemini_response_with_function_calls
assert mock_client.capture.call_count == 1
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
assert call_args["distinct_id"] == "test-id"
assert call_args["event"] == "$ai_generation"
assert props["$ai_provider"] == "gemini"
assert props["$ai_model"] == "gemini-2.5-flash"
assert props["$ai_output_choices"] == [
{
"role": "assistant",
"content": [
{"type": "text", "text": "I'll check the weather for you."},
{"type": "text", "text": " Let me look that up."},
{
"type": "function",
"function": {
"name": "get_current_weather",
"arguments": {"location": "San Francisco"},
},
},
],
}
]
# Check token usage
assert props["$ai_input_tokens"] == 25
assert props["$ai_output_tokens"] == 15
assert props["$ai_http_status"] == 200
async def test_async_cache_and_reasoning_tokens(mock_client, mock_google_genai_client):
"""Test that cache and reasoning tokens are properly extracted with async"""
# Create a mock response with cache and reasoning tokens
mock_response = MagicMock()
mock_response.text = "Test response with cache"
mock_usage = MagicMock()
mock_usage.prompt_token_count = 100
mock_usage.candidates_token_count = 50
mock_usage.cached_content_token_count = 30 # Cache tokens
mock_usage.thoughts_token_count = 10 # Reasoning tokens
mock_response.usage_metadata = mock_usage
# Mock candidates
mock_candidate = MagicMock()
mock_candidate.text = "Test response with cache"
mock_response.candidates = [mock_candidate]
mock_google_genai_client.aio.models.generate_content = AsyncMock(
return_value=mock_response
)
client = AsyncClient(api_key="test-key", insights_client=mock_client)
response = await client.models.generate_content(
model="gemini-2.5-pro",
contents="Test with cache",
insights_distinct_id="test-id",
)
assert response == mock_response
assert mock_client.capture.call_count == 1
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
# Check that all token types are present
assert props["$ai_input_tokens"] == 100
assert props["$ai_output_tokens"] == 50
assert props["$ai_cache_read_input_tokens"] == 30
assert props["$ai_reasoning_tokens"] == 10
async def test_async_streaming_cache_and_reasoning_tokens(
mock_client, mock_google_genai_client
):
"""Test that cache and reasoning tokens are properly extracted in async streaming"""
async def mock_streaming_response():
# Create mock chunks with cache and reasoning tokens
chunk1 = MagicMock()
chunk1.text = "Hello "
chunk1_usage = MagicMock()
chunk1_usage.prompt_token_count = 100
chunk1_usage.candidates_token_count = 5
chunk1_usage.cached_content_token_count = 30 # Cache tokens
chunk1_usage.thoughts_token_count = 0
chunk1.usage_metadata = chunk1_usage
yield chunk1
chunk2 = MagicMock()
chunk2.text = "world!"
chunk2_usage = MagicMock()
chunk2_usage.prompt_token_count = 100
chunk2_usage.candidates_token_count = 10
chunk2_usage.cached_content_token_count = 30 # Same cache tokens
chunk2_usage.thoughts_token_count = 5 # Reasoning tokens
chunk2.usage_metadata = chunk2_usage
yield chunk2
mock_google_genai_client.aio.models.generate_content_stream = AsyncMock(
return_value=mock_streaming_response()
)
client = AsyncClient(api_key="test-key", insights_client=mock_client)
response = await client.models.generate_content_stream(
model="gemini-2.5-pro",
contents="Test streaming with cache",
insights_distinct_id="test-id",
)
# Consume the stream
result = []
async for chunk in response:
result.append(chunk)
assert len(result) == 2
# Check Insights capture was called
assert mock_client.capture.call_count == 1
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
# Check that all token types are present (should use final chunk's usage)
assert props["$ai_input_tokens"] == 100
assert props["$ai_output_tokens"] == 10
assert props["$ai_cache_read_input_tokens"] == 30
assert props["$ai_reasoning_tokens"] == 5
async def test_async_web_search_grounding(mock_client, mock_google_genai_client):
"""Test async web search detection via grounding_metadata."""
# Create mock response with grounding metadata
mock_response = MagicMock()
# Mock usage metadata
mock_usage = MagicMock()
mock_usage.prompt_token_count = 60
mock_usage.candidates_token_count = 40
mock_usage.cached_content_token_count = 0
mock_usage.thoughts_token_count = 0
mock_response.usage_metadata = mock_usage
# Mock grounding metadata
mock_grounding_chunk = MagicMock()
mock_grounding_chunk.uri = "https://example.com"
mock_grounding_metadata = MagicMock()
mock_grounding_metadata.grounding_chunks = [mock_grounding_chunk]
# Mock text part
mock_text_part = MagicMock()
mock_text_part.text = "According to search results..."
type(mock_text_part).text = mock_text_part.text
# Mock content with parts
mock_content = MagicMock()
mock_content.parts = [mock_text_part]
# Mock candidate with grounding metadata
mock_candidate = MagicMock()
mock_candidate.content = mock_content
mock_candidate.grounding_metadata = mock_grounding_metadata
type(mock_candidate).grounding_metadata = mock_candidate.grounding_metadata
mock_response.candidates = [mock_candidate]
mock_response.text = "According to search results..."
# Mock the async generate_content method
mock_google_genai_client.aio.models.generate_content = AsyncMock(
return_value=mock_response
)
client = AsyncClient(api_key="test-key", insights_client=mock_client)
response = await client.models.generate_content(
model="gemini-2.5-flash",
contents="What's the latest news?",
insights_distinct_id="test-id",
)
assert response == mock_response
assert mock_client.capture.call_count == 1
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
# Verify web search count is detected (binary for grounding)
assert props["$ai_web_search_count"] == 1
assert props["$ai_input_tokens"] == 60
assert props["$ai_output_tokens"] == 40
async def test_async_streaming_with_web_search(mock_client, mock_google_genai_client):
"""Test that web search count is properly captured in async streaming mode."""
async def mock_streaming_response():
# Create chunk 1 with grounding metadata
mock_chunk1 = MagicMock()
mock_chunk1.text = "According to "
mock_usage1 = MagicMock()
mock_usage1.prompt_token_count = 30
mock_usage1.candidates_token_count = 5
mock_usage1.cached_content_token_count = 0
mock_usage1.thoughts_token_count = 0
mock_chunk1.usage_metadata = mock_usage1
# Add grounding metadata to first chunk
mock_grounding_chunk = MagicMock()
mock_grounding_chunk.uri = "https://example.com"
mock_grounding_metadata = MagicMock()
mock_grounding_metadata.grounding_chunks = [mock_grounding_chunk]
mock_candidate1 = MagicMock()
mock_candidate1.grounding_metadata = mock_grounding_metadata
type(mock_candidate1).grounding_metadata = mock_candidate1.grounding_metadata
mock_chunk1.candidates = [mock_candidate1]
yield mock_chunk1
# Create chunk 2
mock_chunk2 = MagicMock()
mock_chunk2.text = "search results..."
mock_usage2 = MagicMock()
mock_usage2.prompt_token_count = 30
mock_usage2.candidates_token_count = 15
mock_usage2.cached_content_token_count = 0
mock_usage2.thoughts_token_count = 0
mock_chunk2.usage_metadata = mock_usage2
mock_candidate2 = MagicMock()
mock_chunk2.candidates = [mock_candidate2]
yield mock_chunk2
# Mock the async generate_content_stream method
mock_google_genai_client.aio.models.generate_content_stream = AsyncMock(
return_value=mock_streaming_response()
)
client = AsyncClient(api_key="test-key", insights_client=mock_client)
response = await client.models.generate_content_stream(
model="gemini-2.5-flash",
contents="What's the latest news?",
insights_distinct_id="test-id",
)
chunks = []
async for chunk in response:
chunks.append(chunk)
assert len(chunks) == 2
assert mock_client.capture.call_count == 1
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
# Verify web search count is detected (binary for grounding)
assert props["$ai_web_search_count"] == 1
assert props["$ai_input_tokens"] == 30
assert props["$ai_output_tokens"] == 15
@@ -1,5 +1,5 @@
import pytest
pytest.importorskip("langchain")
pytest.importorskip("langchain_core")
pytest.importorskip("langchain_community")
pytest.importorskip("langgraph")
@@ -21,8 +21,8 @@ try:
from langgraph.graph.state import END, START, StateGraph
from langgraph.prebuilt import create_react_agent
from posthog.ai.langchain import CallbackHandler
from posthog.ai.langchain.callbacks import GenerationMetadata, SpanMetadata
from hanzo_insights.ai.langchain import CallbackHandler
from hanzo_insights.ai.langchain.callbacks import GenerationMetadata, SpanMetadata
LANGCHAIN_AVAILABLE = True
except ImportError:
@@ -53,9 +53,9 @@ ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY")
@pytest.fixture(scope="function")
def mock_client():
with patch("posthog.client.Client") as mock_client:
with patch("hanzo_insights.client.Client") as mock_client:
mock_client.privacy_mode = False
logging.getLogger("posthog").setLevel(logging.DEBUG)
logging.getLogger("hanzo_insights").setLevel(logging.DEBUG)
yield mock_client
@@ -101,7 +101,7 @@ def test_metadata_capture(mock_client):
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"},
metadata={"ls_model_name": "hog-mini", "ls_provider": "hanzo_insights"},
name="test",
)
expected = GenerationMetadata(
@@ -109,10 +109,11 @@ def test_metadata_capture(mock_client):
input=[{"role": "user", "content": "Who won the world series in 2020?"}],
start_time=1234567890,
model_params={"temperature": 0.5},
provider="posthog",
provider="hanzo_insights",
base_url="https://us.posthog.com",
name="test",
end_time=None,
insights_properties=None,
)
assert callbacks._runs[run_id] == expected
with patch("time.time", return_value=1234567891):
@@ -1048,7 +1049,7 @@ def test_base_url_retrieval(mock_client):
prompt = ChatPromptTemplate.from_messages([("user", "Foo")])
chain = prompt | ChatOpenAI(
api_key="test",
model="posthog-mini",
model="insights-mini",
base_url="https://test.posthog.com",
)
callbacks = CallbackHandler(mock_client)
@@ -1124,9 +1125,9 @@ def test_anthropic_chain(mock_client):
)
chain = prompt | ChatAnthropic(
api_key=ANTHROPIC_API_KEY,
model="claude-3-opus-20240229",
model="claude-sonnet-4-5-20250929",
temperature=0,
max_tokens=1,
max_tokens=1024,
)
callbacks = CallbackHandler(
mock_client,
@@ -1149,12 +1150,12 @@ def test_anthropic_chain(mock_client):
assert gen_args["event"] == "$ai_generation"
assert gen_props["$ai_trace_id"] == "test-trace-id"
assert gen_props["$ai_provider"] == "anthropic"
assert gen_props["$ai_model"] == "claude-3-opus-20240229"
assert gen_props["$ai_model"] == "claude-sonnet-4-5-20250929"
assert gen_props["foo"] == "bar"
assert gen_props["$ai_model_parameters"] == {
"temperature": 0.0,
"max_tokens": 1,
"max_tokens": 1024,
"streaming": False,
}
assert gen_props["$ai_input"] == [
@@ -1170,7 +1171,7 @@ def test_anthropic_chain(mock_client):
<= approximate_latency
)
assert gen_props["$ai_input_tokens"] == 17
assert gen_props["$ai_output_tokens"] == 1
assert gen_props["$ai_output_tokens"] == 4
assert trace_args["event"] == "$ai_trace"
assert trace_props["$ai_input_state"] == {}
@@ -1187,9 +1188,9 @@ async def test_async_anthropic_streaming(mock_client):
)
chain = prompt | ChatAnthropic(
api_key=ANTHROPIC_API_KEY,
model="claude-3-opus-20240229",
model="claude-sonnet-4-5-20250929",
temperature=0,
max_tokens=1,
max_tokens=1024,
streaming=True,
stream_usage=True,
)
@@ -1256,7 +1257,7 @@ def test_metadata_tools(mock_client):
run_id,
messages=[{"role": "user", "content": "What's the weather like in SF?"}],
invocation_params={"temperature": 0.5, "tools": tools},
metadata={"ls_model_name": "hog-mini", "ls_provider": "posthog"},
metadata={"ls_model_name": "hog-mini", "ls_provider": "hanzo_insights"},
name="test",
)
expected = GenerationMetadata(
@@ -1264,11 +1265,12 @@ def test_metadata_tools(mock_client):
input=[{"role": "user", "content": "What's the weather like in SF?"}],
start_time=1234567890,
model_params={"temperature": 0.5},
provider="posthog",
provider="hanzo_insights",
base_url="https://us.posthog.com",
name="test",
tools=tools,
end_time=None,
insights_properties=None,
)
assert callbacks._runs[run_id] == expected
with patch("time.time", return_value=1234567891):
@@ -1584,13 +1586,147 @@ def test_anthropic_cache_write_and_read_tokens(mock_client):
generation_props = generation_args["properties"]
assert generation_args["event"] == "$ai_generation"
assert generation_props["$ai_input_tokens"] == 400
assert (
generation_props["$ai_input_tokens"] == 1200
) # No provider metadata, no subtraction
assert generation_props["$ai_output_tokens"] == 30
assert generation_props["$ai_cache_creation_input_tokens"] == 0
assert generation_props["$ai_cache_read_input_tokens"] == 800
assert generation_props["$ai_reasoning_tokens"] == 0
def test_anthropic_provider_subtracts_cache_tokens(mock_client):
"""Test that Anthropic provider correctly subtracts cache tokens from input tokens."""
from langchain_core.outputs import LLMResult, ChatGeneration
from langchain_core.messages import AIMessage
from uuid import uuid4
cb = CallbackHandler(mock_client)
run_id = uuid4()
# Set up with Anthropic provider
cb._set_llm_metadata(
serialized={},
run_id=run_id,
messages=[{"role": "user", "content": "test"}],
metadata={"ls_provider": "anthropic", "ls_model_name": "claude-3-sonnet"},
)
# Response with cache tokens: 1200 input (includes 800 cached)
response = LLMResult(
generations=[
[
ChatGeneration(
message=AIMessage(content="Response"),
generation_info={
"usage_metadata": {
"input_tokens": 1200,
"output_tokens": 50,
"cache_read_input_tokens": 800,
}
},
)
]
],
llm_output={},
)
cb._pop_run_and_capture_generation(run_id, None, response)
generation_args = mock_client.capture.call_args_list[0][1]
assert generation_args["properties"]["$ai_input_tokens"] == 400 # 1200 - 800
assert generation_args["properties"]["$ai_cache_read_input_tokens"] == 800
def test_anthropic_provider_subtracts_cache_write_tokens(mock_client):
"""Test that Anthropic provider correctly subtracts cache write tokens from input tokens."""
from langchain_core.outputs import LLMResult, ChatGeneration
from langchain_core.messages import AIMessage
from uuid import uuid4
cb = CallbackHandler(mock_client)
run_id = uuid4()
# Set up with Anthropic provider
cb._set_llm_metadata(
serialized={},
run_id=run_id,
messages=[{"role": "user", "content": "test"}],
metadata={"ls_provider": "anthropic", "ls_model_name": "claude-3-sonnet"},
)
# Response with cache creation: 1000 input (includes 800 being written to cache)
response = LLMResult(
generations=[
[
ChatGeneration(
message=AIMessage(content="Response"),
generation_info={
"usage_metadata": {
"input_tokens": 1000,
"output_tokens": 50,
"cache_creation_input_tokens": 800,
}
},
)
]
],
llm_output={},
)
cb._pop_run_and_capture_generation(run_id, None, response)
generation_args = mock_client.capture.call_args_list[0][1]
assert generation_args["properties"]["$ai_input_tokens"] == 200 # 1000 - 800
assert generation_args["properties"]["$ai_cache_creation_input_tokens"] == 800
def test_anthropic_provider_subtracts_both_cache_read_and_write_tokens(mock_client):
"""Test that Anthropic provider correctly subtracts both cache read and write tokens."""
from langchain_core.outputs import LLMResult, ChatGeneration
from langchain_core.messages import AIMessage
from uuid import uuid4
cb = CallbackHandler(mock_client)
run_id = uuid4()
# Set up with Anthropic provider
cb._set_llm_metadata(
serialized={},
run_id=run_id,
messages=[{"role": "user", "content": "test"}],
metadata={"ls_provider": "anthropic", "ls_model_name": "claude-3-sonnet"},
)
# Response with both cache read and creation
response = LLMResult(
generations=[
[
ChatGeneration(
message=AIMessage(content="Response"),
generation_info={
"usage_metadata": {
"input_tokens": 2000,
"output_tokens": 50,
"cache_read_input_tokens": 800,
"cache_creation_input_tokens": 500,
}
},
)
]
],
llm_output={},
)
cb._pop_run_and_capture_generation(run_id, None, response)
generation_args = mock_client.capture.call_args_list[0][1]
# 2000 - 800 (read) - 500 (write) = 700
assert generation_args["properties"]["$ai_input_tokens"] == 700
assert generation_args["properties"]["$ai_cache_read_input_tokens"] == 800
assert generation_args["properties"]["$ai_cache_creation_input_tokens"] == 500
def test_openai_cache_read_tokens(mock_client):
"""Test that OpenAI cache read tokens are captured correctly."""
prompt = ChatPromptTemplate.from_messages(
@@ -1626,7 +1762,7 @@ def test_openai_cache_read_tokens(mock_client):
generation_props = generation_args["properties"]
assert generation_args["event"] == "$ai_generation"
assert generation_props["$ai_input_tokens"] == 50
assert generation_props["$ai_input_tokens"] == 150 # No subtraction for OpenAI
assert generation_props["$ai_output_tokens"] == 40
assert generation_props["$ai_cache_read_input_tokens"] == 100
assert generation_props["$ai_cache_creation_input_tokens"] == 0
@@ -1708,7 +1844,7 @@ def test_combined_reasoning_and_cache_tokens(mock_client):
generation_props = generation_args["properties"]
assert generation_args["event"] == "$ai_generation"
assert generation_props["$ai_input_tokens"] == 200
assert generation_props["$ai_input_tokens"] == 500 # No subtraction for OpenAI
assert generation_props["$ai_output_tokens"] == 100
assert generation_props["$ai_cache_read_input_tokens"] == 300
assert generation_props["$ai_cache_creation_input_tokens"] == 0
@@ -1716,7 +1852,7 @@ def test_combined_reasoning_and_cache_tokens(mock_client):
@pytest.mark.skipif(not OPENAI_API_KEY, reason="OPENAI_API_KEY is not set")
def test_openai_reasoning_tokens(mock_client):
def test_openai_reasoning_tokens_o4_mini(mock_client):
model = ChatOpenAI(
api_key=OPENAI_API_KEY, model="o4-mini", max_completion_tokens=10
)
@@ -1731,8 +1867,8 @@ def test_openai_reasoning_tokens(mock_client):
def test_callback_handler_without_client():
"""Test that CallbackHandler works properly when no PostHog client is passed."""
with patch("posthog.ai.langchain.callbacks.setup") as mock_setup:
"""Test that CallbackHandler works properly when no Insights client is passed."""
with patch("hanzo_insights.ai.langchain.callbacks.setup") as mock_setup:
mock_client = mock_setup.return_value
callbacks = CallbackHandler()
@@ -1758,7 +1894,7 @@ def test_callback_handler_without_client():
def test_convert_message_to_dict_tool_calls():
"""Test that _convert_message_to_dict properly converts tool calls in AIMessage."""
from posthog.ai.langchain.callbacks import _convert_message_to_dict
from hanzo_insights.ai.langchain.callbacks import _convert_message_to_dict
from langchain_core.messages import AIMessage
from langchain_core.messages.tool import ToolCall
@@ -1848,7 +1984,7 @@ def test_tool_definition(mock_client):
assert run == expected
assert callbacks._runs == {}
# Now test that the tools are properly captured in the PostHog event
# Now test that the tools are properly captured in the Insights event
mock_response = MagicMock()
mock_response.generations = [[MagicMock()]]
@@ -1917,8 +2053,8 @@ def test_cache_read_tokens_subtraction_from_input_tokens(mock_client):
generation_props = generation_args["properties"]
assert generation_args["event"] == "$ai_generation"
# Input tokens should be reduced: 150 - 100 = 50
assert generation_props["$ai_input_tokens"] == 50
# Input tokens not reduced without provider metadata
assert generation_props["$ai_input_tokens"] == 150
assert generation_props["$ai_output_tokens"] == 40
assert generation_props["$ai_cache_read_input_tokens"] == 100
@@ -1959,8 +2095,8 @@ def test_cache_read_tokens_subtraction_prevents_negative(mock_client):
generation_props = generation_args["properties"]
assert generation_args["event"] == "$ai_generation"
# Input tokens should be 0, not negative: max(80 - 100, 0) = 0
assert generation_props["$ai_input_tokens"] == 0
# Input tokens not reduced without provider metadata
assert generation_props["$ai_input_tokens"] == 80
assert generation_props["$ai_output_tokens"] == 20
assert generation_props["$ai_cache_read_input_tokens"] == 100
@@ -2045,10 +2181,12 @@ def test_zero_input_tokens_with_cache_read(mock_client):
assert generation_props["$ai_cache_read_input_tokens"] == 50
def test_cache_write_tokens_not_subtracted_from_input(mock_client):
"""Test that cache_creation_input_tokens (cache write) do NOT affect input_tokens.
def test_non_anthropic_cache_write_tokens_not_subtracted_from_input(mock_client):
"""Test that cache_creation_input_tokens do NOT affect input_tokens for non-Anthropic providers.
Only cache_read_tokens should be subtracted from input_tokens, not cache_write_tokens.
When no provider metadata is set (or for non-Anthropic providers), cache tokens should
NOT be subtracted from input_tokens. This is because different providers report tokens
differently - only Anthropic's LangChain integration requires subtraction.
"""
prompt = ChatPromptTemplate.from_messages([("user", "Create cache")])
@@ -2098,7 +2236,7 @@ def test_agent_action_and_finish_imports():
from langchain.schema.agent import AgentAction, AgentFinish # type: ignore
# Verify they're available in the callbacks module
from posthog.ai.langchain.callbacks import CallbackHandler
from hanzo_insights.ai.langchain.callbacks import CallbackHandler
# Test on_agent_action with mock data
mock_client = MagicMock()
@@ -2126,3 +2264,383 @@ def test_agent_action_and_finish_imports():
assert mock_client.capture.call_count == 1
call_args = mock_client.capture.call_args[1]
assert call_args["event"] == "$ai_span"
def test_insights_properties_field_in_generation_metadata(mock_client):
"""Test that insights_properties is properly stored in GenerationMetadata."""
callbacks = CallbackHandler(mock_client)
run_id = uuid.uuid4()
# Test with billable=True
with patch("time.time", return_value=1234567890):
callbacks._set_llm_metadata(
{"kwargs": {"openai_api_base": "https://api.openai.com"}},
run_id,
messages=[{"role": "user", "content": "Test message"}],
invocation_params={"temperature": 0.5},
metadata={
"ls_model_name": "gpt-4o",
"ls_provider": "openai",
"insights_properties": {"$ai_billable": True},
},
name="test",
)
expected = GenerationMetadata(
model="gpt-4o",
input=[{"role": "user", "content": "Test message"}],
start_time=1234567890,
model_params={"temperature": 0.5},
provider="openai",
base_url="https://api.openai.com",
name="test",
insights_properties={"$ai_billable": True},
end_time=None,
)
assert callbacks._runs[run_id] == expected
assert callbacks._runs[run_id].insights_properties == {"$ai_billable": True}
callbacks._pop_run_metadata(run_id)
# Test with billable=False (explicit)
run_id2 = uuid.uuid4()
with patch("time.time", return_value=1234567890):
callbacks._set_llm_metadata(
{"kwargs": {"openai_api_base": "https://api.openai.com"}},
run_id2,
messages=[{"role": "user", "content": "Test message"}],
invocation_params={"temperature": 0.5},
metadata={
"ls_model_name": "gpt-4o",
"ls_provider": "openai",
"insights_properties": {"$ai_billable": False},
},
name="test",
)
assert callbacks._runs[run_id2].insights_properties == {"$ai_billable": False}
callbacks._pop_run_metadata(run_id2)
# Test when insights_properties not provided
run_id3 = uuid.uuid4()
with patch("time.time", return_value=1234567890):
callbacks._set_llm_metadata(
{"kwargs": {"openai_api_base": "https://api.openai.com"}},
run_id3,
messages=[{"role": "user", "content": "Test message"}],
invocation_params={"temperature": 0.5},
metadata={"ls_model_name": "gpt-4o", "ls_provider": "openai"},
name="test",
)
assert callbacks._runs[run_id3].insights_properties is None
def test_billable_property_in_generation_event(mock_client):
"""Test that the billable property is captured in the $ai_generation event."""
callbacks = CallbackHandler(mock_client)
# We need to test the _set_llm_metadata directly since FakeMessagesListChatModel
# doesn't support metadata in the same way as real models
run_id = uuid.uuid4()
with patch("time.time", return_value=1234567890):
callbacks._set_llm_metadata(
{},
run_id,
messages=[{"role": "user", "content": "Test"}],
metadata={
"insights_properties": {"$ai_billable": True},
"ls_model_name": "test-model",
},
invocation_params={},
)
mock_response = MagicMock()
mock_response.generations = [[MagicMock()]]
with patch("time.time", return_value=1234567891):
run = callbacks._pop_run_metadata(run_id)
callbacks._capture_generation(
trace_id=run_id,
run_id=run_id,
run=run,
output=mock_response,
parent_run_id=None,
)
assert mock_client.capture.call_count == 1
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
assert call_args["event"] == "$ai_generation"
assert props["$ai_billable"] is True
def test_billable_defaults_to_false_in_event(mock_client):
"""Test that $ai_billable is not present when not specified."""
prompt = ChatPromptTemplate.from_messages([("user", "Test query")])
model = FakeMessagesListChatModel(
responses=[AIMessage(content="Test response")],
)
callbacks = [CallbackHandler(mock_client)]
chain = prompt | model
chain.invoke({}, config={"callbacks": callbacks})
generation_call = None
for call in mock_client.capture.call_args_list:
if call[1]["event"] == "$ai_generation":
generation_call = call
break
assert generation_call is not None
props = generation_call[1]["properties"]
assert "$ai_billable" not in props
def test_billable_with_real_chain(mock_client):
"""Test billable tracking through a complete chain execution with mocked metadata."""
callbacks = CallbackHandler(mock_client)
run_id = uuid.uuid4()
with patch("time.time", return_value=1000.0):
callbacks._set_llm_metadata(
{},
run_id,
messages=[{"role": "user", "content": "What's the weather?"}],
metadata={
"ls_model_name": "fake-model",
"ls_provider": "fake",
"insights_properties": {"$ai_billable": True},
},
invocation_params={"temperature": 0.7},
)
assert callbacks._runs[run_id].insights_properties == {"$ai_billable": True}
mock_response = MagicMock()
mock_response.generations = [[MagicMock()]]
with patch("time.time", return_value=1001.0):
run = callbacks._pop_run_metadata(run_id)
callbacks._capture_generation(
trace_id=run_id,
run_id=run_id,
run=run,
output=mock_response,
parent_run_id=None,
)
assert mock_client.capture.call_count == 1
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
assert call_args["event"] == "$ai_generation"
assert props["$ai_billable"] is True
assert props["$ai_model"] == "fake-model"
assert props["$ai_provider"] == "fake"
# Exception Capture Integration Tests
def test_exception_autocapture_on_span_error():
"""Test that capture_exception is called when a span errors and autocapture is enabled."""
mock_client = MagicMock()
mock_client.privacy_mode = False
mock_client.enable_exception_autocapture = True
mock_client.capture_exception.return_value = "exception-uuid-123"
def failing_span(_):
raise ValueError("test error")
callbacks = [CallbackHandler(mock_client)]
chain = RunnableLambda(failing_span)
try:
chain.invoke({}, config={"callbacks": callbacks})
except ValueError:
pass
# Verify capture_exception was called
assert mock_client.capture_exception.call_count == 1
exception_call = mock_client.capture_exception.call_args
assert isinstance(exception_call[0][0], ValueError)
assert str(exception_call[0][0]) == "test error"
def test_exception_autocapture_adds_exception_id_to_span_event():
"""Test that $exception_event_id is added to the span event properties."""
mock_client = MagicMock()
mock_client.privacy_mode = False
mock_client.enable_exception_autocapture = True
mock_client.capture_exception.return_value = "exception-uuid-456"
def failing_span(_):
raise ValueError("test error")
callbacks = [CallbackHandler(mock_client)]
chain = RunnableLambda(failing_span)
try:
chain.invoke({}, config={"callbacks": callbacks})
except ValueError:
pass
# Find the span event (should have $ai_is_error=True)
span_calls = [
call
for call in mock_client.capture.call_args_list
if call[1].get("properties", {}).get("$ai_is_error") is True
]
assert len(span_calls) >= 1
span_props = span_calls[0][1]["properties"]
assert span_props["$exception_event_id"] == "exception-uuid-456"
assert span_props["$ai_error"] == "ValueError: test error"
def test_exception_autocapture_disabled_does_not_capture():
"""Test that capture_exception is NOT called when autocapture is disabled."""
mock_client = MagicMock()
mock_client.privacy_mode = False
mock_client.enable_exception_autocapture = False
def failing_span(_):
raise ValueError("test error")
callbacks = [CallbackHandler(mock_client)]
chain = RunnableLambda(failing_span)
try:
chain.invoke({}, config={"callbacks": callbacks})
except ValueError:
pass
# Verify capture_exception was NOT called
assert mock_client.capture_exception.call_count == 0
# But the span event should still have error info
span_calls = [
call
for call in mock_client.capture.call_args_list
if call[1].get("properties", {}).get("$ai_is_error") is True
]
assert len(span_calls) >= 1
span_props = span_calls[0][1]["properties"]
assert "$exception_event_id" not in span_props
assert span_props["$ai_error"] == "ValueError: test error"
def test_exception_autocapture_on_llm_generation_error(mock_client):
"""Test that capture_exception is called when an LLM generation fails."""
mock_client.privacy_mode = False
mock_client.enable_exception_autocapture = True
mock_client.capture_exception.return_value = "exception-uuid-789"
callbacks = CallbackHandler(mock_client)
run_id = uuid.uuid4()
# Simulate LLM start
callbacks.on_llm_start(
serialized={"kwargs": {"openai_api_base": "https://api.openai.com"}},
prompts=["Hello"],
run_id=run_id,
)
# Simulate LLM error
error = Exception("API rate limit exceeded")
callbacks.on_llm_error(error, run_id=run_id)
# Verify capture_exception was called
assert mock_client.capture_exception.call_count == 1
exception_call = mock_client.capture_exception.call_args
assert exception_call[0][0] is error
# Verify the generation event has $exception_event_id
generation_calls = [
call
for call in mock_client.capture.call_args_list
if call[1].get("event") == "$ai_generation"
]
assert len(generation_calls) == 1
gen_props = generation_calls[0][1]["properties"]
assert gen_props["$exception_event_id"] == "exception-uuid-789"
assert gen_props["$ai_is_error"] is True
def test_exception_autocapture_passes_ai_properties_to_exception():
"""Test that AI properties are passed to the exception event."""
mock_client = MagicMock()
mock_client.privacy_mode = False
mock_client.enable_exception_autocapture = True
mock_client.capture_exception.return_value = "exception-uuid-abc"
callbacks = CallbackHandler(
mock_client,
distinct_id="user-123",
properties={"custom_prop": "custom_value"},
)
run_id = uuid.uuid4()
# Simulate LLM start
callbacks.on_llm_start(
serialized={"kwargs": {"openai_api_base": "https://api.openai.com"}},
prompts=["Hello"],
run_id=run_id,
)
# Simulate LLM error
error = Exception("API error")
callbacks.on_llm_error(error, run_id=run_id)
# Verify capture_exception received the properties
exception_call = mock_client.capture_exception.call_args
props = exception_call[1]["properties"]
# Should have AI-related properties
assert "$ai_trace_id" in props
assert "$ai_is_error" in props
assert props["$ai_is_error"] is True
# Should have distinct_id passed through
assert exception_call[1]["distinct_id"] == "user-123"
def test_exception_autocapture_none_return_no_exception_id():
"""Test that when capture_exception returns None, no $exception_event_id is added."""
mock_client = MagicMock()
mock_client.privacy_mode = False
mock_client.enable_exception_autocapture = True
mock_client.capture_exception.return_value = (
None # e.g., exception already captured
)
def failing_span(_):
raise ValueError("test error")
callbacks = [CallbackHandler(mock_client)]
chain = RunnableLambda(failing_span)
try:
chain.invoke({}, config={"callbacks": callbacks})
except ValueError:
pass
# capture_exception was called but returned None
assert mock_client.capture_exception.call_count == 1
# Span event should NOT have $exception_event_id
span_calls = [
call
for call in mock_client.capture.call_args_list
if call[1].get("properties", {}).get("$ai_is_error") is True
]
assert len(span_calls) >= 1
span_props = span_calls[0][1]["properties"]
assert "$exception_event_id" not in span_props
@@ -1,3 +1,4 @@
import json
import time
from unittest.mock import AsyncMock, patch
@@ -33,8 +34,8 @@ try:
ParsedResponseOutputText,
)
from posthog.ai.openai import OpenAI
from posthog.ai.openai.openai_async import AsyncOpenAI
from hanzo_insights.ai.openai import OpenAI
from hanzo_insights.ai.openai.openai_async import AsyncOpenAI
OPENAI_AVAILABLE = True
except ImportError:
@@ -48,7 +49,7 @@ pytestmark = pytest.mark.skipif(
@pytest.fixture
def mock_client():
with patch("posthog.client.Client") as mock_client:
with patch("hanzo_insights.client.Client") as mock_client:
mock_client.privacy_mode = False
yield mock_client
@@ -466,12 +467,12 @@ def test_basic_completion(mock_client, mock_openai_response):
"openai.resources.chat.completions.Completions.create",
return_value=mock_openai_response,
):
client = OpenAI(api_key="test-key", posthog_client=mock_client)
client = OpenAI(api_key="test-key", insights_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"},
insights_distinct_id="test-id",
insights_properties={"foo": "bar"},
)
assert response == mock_openai_response
@@ -496,6 +497,15 @@ def test_basic_completion(mock_client, mock_openai_response):
assert props["$ai_http_status"] == 200
assert props["foo"] == "bar"
assert isinstance(props["$ai_latency"], float)
# Verify raw usage metadata is passed for backend processing
assert "$ai_usage" in props
assert props["$ai_usage"] is not None
# Verify it's JSON-serializable
json.dumps(props["$ai_usage"])
# Verify it has expected structure
assert isinstance(props["$ai_usage"], dict)
assert "prompt_tokens" in props["$ai_usage"]
assert "completion_tokens" in props["$ai_usage"]
def test_embeddings(mock_client, mock_embedding_response):
@@ -503,12 +513,12 @@ def test_embeddings(mock_client, mock_embedding_response):
"openai.resources.embeddings.Embeddings.create",
return_value=mock_embedding_response,
):
client = OpenAI(api_key="test-key", posthog_client=mock_client)
client = OpenAI(api_key="test-key", insights_client=mock_client)
response = client.embeddings.create(
model="text-embedding-3-small",
input="Hello world",
posthog_distinct_id="test-id",
posthog_properties={"foo": "bar"},
insights_distinct_id="test-id",
insights_properties={"foo": "bar"},
)
assert response == mock_embedding_response
@@ -533,12 +543,12 @@ def test_groups(mock_client, mock_openai_response):
"openai.resources.chat.completions.Completions.create",
return_value=mock_openai_response,
):
client = OpenAI(api_key="test-key", posthog_client=mock_client)
client = OpenAI(api_key="test-key", insights_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"},
insights_distinct_id="test-id",
insights_groups={"company": "test_company"},
)
assert response == mock_openai_response
@@ -554,12 +564,12 @@ def test_privacy_mode_local(mock_client, mock_openai_response):
"openai.resources.chat.completions.Completions.create",
return_value=mock_openai_response,
):
client = OpenAI(api_key="test-key", posthog_client=mock_client)
client = OpenAI(api_key="test-key", insights_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,
insights_distinct_id="test-id",
insights_privacy_mode=True,
)
assert response == mock_openai_response
@@ -577,12 +587,12 @@ def test_privacy_mode_global(mock_client, mock_openai_response):
return_value=mock_openai_response,
):
mock_client.privacy_mode = True
client = OpenAI(api_key="test-key", posthog_client=mock_client)
client = OpenAI(api_key="test-key", insights_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,
insights_distinct_id="test-id",
insights_privacy_mode=False,
)
assert response == mock_openai_response
@@ -599,7 +609,7 @@ def test_error(mock_client, mock_openai_response):
"openai.resources.chat.completions.Completions.create",
side_effect=Exception("Test error"),
):
client = OpenAI(api_key="test-key", posthog_client=mock_client)
client = OpenAI(api_key="test-key", insights_client=mock_client)
with pytest.raises(Exception):
client.chat.completions.create(
model="gpt-4", messages=[{"role": "user", "content": "Hello"}]
@@ -618,12 +628,12 @@ def test_cached_tokens(mock_client, mock_openai_response_with_cached_tokens):
"openai.resources.chat.completions.Completions.create",
return_value=mock_openai_response_with_cached_tokens,
):
client = OpenAI(api_key="test-key", posthog_client=mock_client)
client = OpenAI(api_key="test-key", insights_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"},
insights_distinct_id="test-id",
insights_properties={"foo": "bar"},
)
assert response == mock_openai_response_with_cached_tokens
@@ -656,7 +666,7 @@ def test_tool_calls(mock_client, mock_openai_response_with_tool_calls):
"openai.resources.chat.completions.Completions.create",
return_value=mock_openai_response_with_tool_calls,
):
client = OpenAI(api_key="test-key", posthog_client=mock_client)
client = OpenAI(api_key="test-key", insights_client=mock_client)
response = client.chat.completions.create(
model="gpt-4",
messages=[
@@ -672,7 +682,7 @@ def test_tool_calls(mock_client, mock_openai_response_with_tool_calls):
},
}
],
posthog_distinct_id="test-id",
insights_distinct_id="test-id",
)
assert response == mock_openai_response_with_tool_calls
@@ -729,7 +739,7 @@ def test_tool_calls_only_no_content(mock_client, mock_openai_response_tool_calls
"openai.resources.chat.completions.Completions.create",
return_value=mock_openai_response_tool_calls_only,
):
client = OpenAI(api_key="test-key", posthog_client=mock_client)
client = OpenAI(api_key="test-key", insights_client=mock_client)
response = client.chat.completions.create(
model="gpt-4",
messages=[{"role": "user", "content": "Get weather for New York"}],
@@ -743,7 +753,7 @@ def test_tool_calls_only_no_content(mock_client, mock_openai_response_tool_calls
},
}
],
posthog_distinct_id="test-id",
insights_distinct_id="test-id",
)
assert response == mock_openai_response_tool_calls_only
@@ -783,7 +793,7 @@ def test_responses_api_tool_calls(mock_client, mock_responses_api_with_tool_call
"openai.resources.responses.Responses.create",
return_value=mock_responses_api_with_tool_calls,
):
client = OpenAI(api_key="test-key", posthog_client=mock_client)
client = OpenAI(api_key="test-key", insights_client=mock_client)
response = client.responses.create(
model="gpt-4o-mini",
input=[{"role": "user", "content": "What's the weather in Chicago?"}],
@@ -798,7 +808,7 @@ def test_responses_api_tool_calls(mock_client, mock_responses_api_with_tool_call
},
}
],
posthog_distinct_id="test-id",
insights_distinct_id="test-id",
)
assert response == mock_responses_api_with_tool_calls
@@ -841,7 +851,7 @@ def test_streaming_with_tool_calls(mock_client, streaming_tool_call_chunks):
# Set up the mock to return our chunks when iterated
mock_create.return_value = streaming_tool_call_chunks
client = OpenAI(api_key="test-key", posthog_client=mock_client)
client = OpenAI(api_key="test-key", insights_client=mock_client)
# Call the streaming method
response_generator = client.chat.completions.create(
@@ -860,7 +870,7 @@ def test_streaming_with_tool_calls(mock_client, streaming_tool_call_chunks):
}
],
stream=True,
posthog_distinct_id="test-id",
insights_distinct_id="test-id",
)
# Consume the generator to trigger the event capture
@@ -922,6 +932,16 @@ def test_streaming_with_tool_calls(mock_client, streaming_tool_call_chunks):
assert props["$ai_input_tokens"] == 20
assert props["$ai_output_tokens"] == 15
# Verify raw usage is captured in streaming mode
assert "$ai_usage" in props
assert props["$ai_usage"] is not None
# Verify it's JSON-serializable
json.dumps(props["$ai_usage"])
# Verify it has expected structure (merged from chunks)
assert isinstance(props["$ai_usage"], dict)
assert "prompt_tokens" in props["$ai_usage"]
assert "completion_tokens" in props["$ai_usage"]
# test responses api
def test_responses_api(mock_client, mock_openai_response_with_responses_api):
@@ -929,12 +949,12 @@ def test_responses_api(mock_client, mock_openai_response_with_responses_api):
"openai.resources.responses.Responses.create",
return_value=mock_openai_response_with_responses_api,
):
client = OpenAI(api_key="test-key", posthog_client=mock_client)
client = OpenAI(api_key="test-key", insights_client=mock_client)
response = client.responses.create(
model="gpt-4o-mini",
input="Hello",
posthog_distinct_id="test-id",
posthog_properties={"foo": "bar"},
insights_distinct_id="test-id",
insights_properties={"foo": "bar"},
)
assert response == mock_openai_response_with_responses_api
assert mock_client.capture.call_count == 1
@@ -966,7 +986,7 @@ def test_responses_parse(mock_client, mock_parsed_response):
"openai.resources.responses.Responses.parse",
return_value=mock_parsed_response,
):
client = OpenAI(api_key="test-key", posthog_client=mock_client)
client = OpenAI(api_key="test-key", insights_client=mock_client)
response = client.responses.parse(
model="gpt-4o-2024-08-06",
input=[
@@ -996,8 +1016,8 @@ def test_responses_parse(mock_client, mock_parsed_response):
},
}
},
posthog_distinct_id="test-id",
posthog_properties={"foo": "bar"},
insights_distinct_id="test-id",
insights_properties={"foo": "bar"},
)
assert response == mock_parsed_response
@@ -1081,15 +1101,15 @@ def test_responses_api_streaming_with_tokens(mock_client):
"openai.resources.responses.Responses.create",
side_effect=mock_streaming_response,
):
client = OpenAI(api_key="test-key", posthog_client=mock_client)
client = OpenAI(api_key="test-key", insights_client=mock_client)
# Consume the streaming response
response = client.responses.create(
model="gpt-4o-mini",
input=[{"role": "user", "content": "Test message"}],
stream=True,
posthog_distinct_id="test-id",
posthog_properties={"test": "streaming"},
insights_distinct_id="test-id",
insights_properties={"test": "streaming"},
)
# Consume all chunks
@@ -1133,7 +1153,7 @@ async def test_async_chat_streaming_with_tool_calls(
with patch(
"openai.resources.chat.completions.AsyncCompletions.create", new=mock_create
):
client = AsyncOpenAI(api_key="test-key", posthog_client=mock_client)
client = AsyncOpenAI(api_key="test-key", insights_client=mock_client)
response_stream = await client.chat.completions.create(
model="gpt-4",
@@ -1151,7 +1171,7 @@ async def test_async_chat_streaming_with_tool_calls(
}
],
stream=True,
posthog_distinct_id="test-id",
insights_distinct_id="test-id",
)
chunks = []
@@ -1219,14 +1239,14 @@ async def test_async_responses_streaming_with_tokens(mock_client):
return chunk_iterable()
with patch("openai.resources.responses.AsyncResponses.create", new=mock_create):
client = AsyncOpenAI(api_key="test-key", posthog_client=mock_client)
client = AsyncOpenAI(api_key="test-key", insights_client=mock_client)
response_stream = await client.responses.create(
model="gpt-4o-mini",
input=[{"role": "user", "content": "Test message"}],
stream=True,
posthog_distinct_id="test-id",
posthog_properties={"test": "streaming"},
insights_distinct_id="test-id",
insights_properties={"test": "streaming"},
)
async for _ in response_stream:
@@ -1254,13 +1274,13 @@ async def test_async_embeddings_create(mock_client, mock_embedding_response):
mock_create = AsyncMock(return_value=mock_embedding_response)
with patch("openai.resources.embeddings.AsyncEmbeddings.create", new=mock_create):
client = AsyncOpenAI(api_key="test-key", posthog_client=mock_client)
client = AsyncOpenAI(api_key="test-key", insights_client=mock_client)
response = await client.embeddings.create(
model="text-embedding-3-small",
input="Hello world",
posthog_distinct_id="test-id",
posthog_properties={"foo": "bar"},
insights_distinct_id="test-id",
insights_properties={"foo": "bar"},
)
assert response == mock_embedding_response
@@ -1284,7 +1304,7 @@ def test_tool_definition(mock_client, mock_openai_response):
"openai.resources.chat.completions.Completions.create",
return_value=mock_openai_response,
):
client = OpenAI(api_key="test-key", posthog_client=mock_client)
client = OpenAI(api_key="test-key", insights_client=mock_client)
# Define tools to be passed to the create function
tools = [
@@ -1311,8 +1331,8 @@ def test_tool_definition(mock_client, mock_openai_response):
model="gpt-4o-mini",
messages=[{"role": "user", "content": "hey"}],
tools=tools,
posthog_distinct_id="test-id",
posthog_properties={"foo": "bar"},
insights_distinct_id="test-id",
insights_properties={"foo": "bar"},
)
assert response == mock_openai_response
@@ -1372,11 +1392,11 @@ def test_web_search_perplexity_style(mock_client):
mock_response = MockResponseWithAnnotations()
with patch("openai.resources.chat.Completions.create", return_value=mock_response):
client = OpenAI(api_key="test-key", posthog_client=mock_client)
client = OpenAI(api_key="test-key", insights_client=mock_client)
response = client.chat.completions.create(
model="gpt-4-turbo",
messages=[{"role": "user", "content": "What's happening in tech?"}],
posthog_distinct_id="test-id",
insights_distinct_id="test-id",
)
assert response == mock_response
@@ -1419,19 +1439,19 @@ def test_web_search_responses_api(mock_client):
"openai.resources.responses.Responses.create", return_value=mock_response
):
# Manually call the tracking since we're testing the converter logic
from posthog.ai.utils import call_llm_and_track_usage
from hanzo_insights.ai.utils import call_llm_and_track_usage
def mock_create_call(**kwargs):
return mock_response
result = call_llm_and_track_usage(
posthog_distinct_id="test-id",
insights_distinct_id="test-id",
ph_client=mock_client,
provider="openai",
posthog_trace_id=None,
posthog_properties=None,
posthog_privacy_mode=False,
posthog_groups=None,
insights_trace_id=None,
insights_properties=None,
insights_privacy_mode=False,
insights_groups=None,
base_url="https://api.openai.com/v1",
call_method=mock_create_call,
model="gpt-4o",
@@ -1513,12 +1533,12 @@ def test_streaming_with_web_search(mock_client, streaming_web_search_chunks):
with patch("openai.resources.chat.completions.Completions.create") as mock_create:
mock_create.return_value = streaming_web_search_chunks
client = OpenAI(api_key="test-key", posthog_client=mock_client)
client = OpenAI(api_key="test-key", insights_client=mock_client)
response_generator = client.chat.completions.create(
model="gpt-4",
messages=[{"role": "user", "content": "Search for recent news"}],
stream=True,
posthog_distinct_id="test-id",
insights_distinct_id="test-id",
)
# Consume the generator to trigger the event capture
@@ -1549,12 +1569,12 @@ def test_streaming_with_web_search_on_non_usage_chunk(
with patch("openai.resources.chat.completions.Completions.create") as mock_create:
mock_create.return_value = streaming_web_search_chunks
client = OpenAI(api_key="test-key", posthog_client=mock_client)
client = OpenAI(api_key="test-key", insights_client=mock_client)
response_generator = client.chat.completions.create(
model="gpt-4",
messages=[{"role": "user", "content": "Search for recent news"}],
stream=True,
posthog_distinct_id="test-id",
insights_distinct_id="test-id",
)
# Consume the generator to trigger the event capture
@@ -1609,12 +1629,12 @@ async def test_async_chat_with_web_search(mock_client):
with patch(
"openai.resources.chat.completions.AsyncCompletions.create", new=mock_create
):
client = AsyncOpenAI(api_key="test-key", posthog_client=mock_client)
client = AsyncOpenAI(api_key="test-key", insights_client=mock_client)
response = await client.chat.completions.create(
model="gpt-4",
messages=[{"role": "user", "content": "Search for recent news"}],
posthog_distinct_id="test-id",
insights_distinct_id="test-id",
)
assert response == mock_response
@@ -1652,13 +1672,13 @@ async def test_async_chat_streaming_with_web_search(
with patch(
"openai.resources.chat.completions.AsyncCompletions.create", new=mock_create
):
client = AsyncOpenAI(api_key="test-key", posthog_client=mock_client)
client = AsyncOpenAI(api_key="test-key", insights_client=mock_client)
response_stream = await client.chat.completions.create(
model="gpt-4",
messages=[{"role": "user", "content": "Search for recent news"}],
stream=True,
posthog_distinct_id="test-id",
insights_distinct_id="test-id",
)
chunks = []
@@ -1676,3 +1696,459 @@ async def test_async_chat_streaming_with_web_search(
assert props["$ai_web_search_count"] == 1
assert props["$ai_input_tokens"] == 20
assert props["$ai_output_tokens"] == 15
# Tests for model extraction fallback (stored prompts support)
def test_streaming_chat_extracts_model_from_chunk_when_not_in_kwargs(mock_client):
"""Test that model is extracted from streaming chunks when not provided in kwargs (stored prompts)."""
# Create streaming chunks with model field but we won't pass model in kwargs
chunks = [
ChatCompletionChunk(
id="chunk1",
model="gpt-4o-stored-prompt", # Model comes from response, not request
object="chat.completion.chunk",
created=1234567890,
choices=[
ChoiceChunk(
index=0,
delta=ChoiceDelta(role="assistant", content="Hello"),
finish_reason=None,
)
],
),
ChatCompletionChunk(
id="chunk2",
model="gpt-4o-stored-prompt",
object="chat.completion.chunk",
created=1234567891,
choices=[
ChoiceChunk(
index=0,
delta=ChoiceDelta(content=" world"),
finish_reason="stop",
)
],
usage=CompletionUsage(
prompt_tokens=10,
completion_tokens=5,
total_tokens=15,
),
),
]
with patch("openai.resources.chat.completions.Completions.create") as mock_create:
mock_create.return_value = chunks
client = OpenAI(api_key="test-key", insights_client=mock_client)
# Note: NOT passing model in kwargs - simulates stored prompt usage
response_generator = client.chat.completions.create(
messages=[{"role": "user", "content": "Hello"}],
stream=True,
insights_distinct_id="test-id",
)
# Consume the generator
list(response_generator)
assert mock_client.capture.call_count == 1
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
# Model should be extracted from chunk, not kwargs
assert props["$ai_model"] == "gpt-4o-stored-prompt"
def test_streaming_chat_prefers_kwargs_model_over_chunk_model(mock_client):
"""Test that model from kwargs takes precedence over model from chunk."""
chunks = [
ChatCompletionChunk(
id="chunk1",
model="gpt-4o-from-response",
object="chat.completion.chunk",
created=1234567890,
choices=[
ChoiceChunk(
index=0,
delta=ChoiceDelta(role="assistant", content="Hello"),
finish_reason="stop",
)
],
usage=CompletionUsage(
prompt_tokens=10,
completion_tokens=5,
total_tokens=15,
),
),
]
with patch("openai.resources.chat.completions.Completions.create") as mock_create:
mock_create.return_value = chunks
client = OpenAI(api_key="test-key", insights_client=mock_client)
response_generator = client.chat.completions.create(
model="gpt-4o-from-kwargs", # Explicitly passed model
messages=[{"role": "user", "content": "Hello"}],
stream=True,
insights_distinct_id="test-id",
)
list(response_generator)
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
# kwargs model should take precedence
assert props["$ai_model"] == "gpt-4o-from-kwargs"
def test_streaming_responses_api_extracts_model_from_response_object(mock_client):
"""Test that Responses API streaming extracts model from chunk.response.model (stored prompts)."""
from unittest.mock import MagicMock
from openai.types.responses import ResponseUsage
chunks = []
# Content chunk
chunk1 = MagicMock()
chunk1.type = "response.text.delta"
chunk1.text = "Test response"
# No response attribute on content chunks
del chunk1.response
chunks.append(chunk1)
# Final chunk with response object containing model
chunk2 = MagicMock()
chunk2.type = "response.completed"
chunk2.response = MagicMock()
chunk2.response.model = "gpt-4o-mini-stored" # Model from stored prompt
chunk2.response.usage = ResponseUsage(
input_tokens=20,
output_tokens=10,
total_tokens=30,
input_tokens_details={"prompt_tokens": 20, "cached_tokens": 0},
output_tokens_details={"reasoning_tokens": 0},
)
chunk2.response.output = ["Test response"]
chunks.append(chunk2)
with patch("openai.resources.responses.Responses.create") as mock_create:
mock_create.return_value = iter(chunks)
client = OpenAI(api_key="test-key", insights_client=mock_client)
# Note: NOT passing model - simulates stored prompt
response_generator = client.responses.create(
input=[{"role": "user", "content": "Hello"}],
stream=True,
insights_distinct_id="test-id",
)
list(response_generator)
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
# Model should be extracted from chunk.response.model
assert props["$ai_model"] == "gpt-4o-mini-stored"
def test_non_streaming_extracts_model_from_response(mock_client):
"""Test that non-streaming calls extract model from response when not in kwargs."""
# Create a response with model but we won't pass model in kwargs
mock_response = ChatCompletion(
id="test",
model="gpt-4o-stored-prompt",
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,
),
)
with patch(
"openai.resources.chat.completions.Completions.create",
return_value=mock_response,
):
client = OpenAI(api_key="test-key", insights_client=mock_client)
# Note: NOT passing model in kwargs
response = client.chat.completions.create(
messages=[{"role": "user", "content": "Hello"}],
insights_distinct_id="test-id",
)
assert response == mock_response
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
# Model should be extracted from response.model
assert props["$ai_model"] == "gpt-4o-stored-prompt"
def test_non_streaming_responses_api_extracts_model_from_response(mock_client):
"""Test that non-streaming Responses API extracts model from response when not in kwargs."""
mock_response = Response(
id="test",
model="gpt-4o-mini-stored",
object="response",
created_at=1741476542,
status="completed",
error=None,
incomplete_details=None,
instructions=None,
max_output_tokens=None,
tools=[],
tool_choice="auto",
output=[
ResponseOutputMessage(
id="msg_123",
type="message",
role="assistant",
status="completed",
content=[
ResponseOutputText(
type="output_text",
text="Test response",
annotations=[],
)
],
)
],
parallel_tool_calls=True,
previous_response_id=None,
usage=ResponseUsage(
input_tokens=10,
output_tokens=10,
input_tokens_details={"prompt_tokens": 10, "cached_tokens": 0},
output_tokens_details={"reasoning_tokens": 0},
total_tokens=20,
),
user=None,
metadata={},
)
with patch(
"openai.resources.responses.Responses.create",
return_value=mock_response,
):
client = OpenAI(api_key="test-key", insights_client=mock_client)
# Note: NOT passing model in kwargs
response = client.responses.create(
input="Hello",
insights_distinct_id="test-id",
)
assert response == mock_response
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
# Model should be extracted from response.model
assert props["$ai_model"] == "gpt-4o-mini-stored"
def test_non_streaming_returns_none_when_no_model(mock_client):
"""Test that non-streaming returns None (not 'unknown') when model is not available anywhere."""
# Create a response without model attribute using real OpenAI types
mock_response = ChatCompletion(
id="test",
model="", # Will be removed below
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=5,
prompt_tokens=10,
total_tokens=15,
),
)
# Remove model attribute to simulate missing model
object.__delattr__(mock_response, "model")
with patch(
"openai.resources.chat.completions.Completions.create",
return_value=mock_response,
):
client = OpenAI(api_key="test-key", insights_client=mock_client)
# Note: NOT passing model in kwargs and response has no model
client.chat.completions.create(
messages=[{"role": "user", "content": "Hello"}],
insights_distinct_id="test-id",
)
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
# Should be None, NOT "unknown" (to avoid incorrect cost matching)
assert props["$ai_model"] is None
def test_streaming_falls_back_to_unknown_when_no_model(mock_client):
"""Test that streaming falls back to 'unknown' when model is not available anywhere."""
from unittest.mock import MagicMock
# Create a chunk without model attribute
chunk = MagicMock()
chunk.choices = [MagicMock()]
chunk.choices[0].delta = MagicMock()
chunk.choices[0].delta.content = "Hello"
chunk.choices[0].delta.role = "assistant"
chunk.choices[0].delta.tool_calls = None
chunk.usage = CompletionUsage(
prompt_tokens=10,
completion_tokens=5,
total_tokens=15,
)
# Explicitly remove model attribute
del chunk.model
with patch("openai.resources.chat.completions.Completions.create") as mock_create:
mock_create.return_value = [chunk]
client = OpenAI(api_key="test-key", insights_client=mock_client)
response_generator = client.chat.completions.create(
messages=[{"role": "user", "content": "Hello"}],
stream=True,
insights_distinct_id="test-id",
)
list(response_generator)
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
# Should fall back to "unknown"
assert props["$ai_model"] == "unknown"
@pytest.mark.asyncio
async def test_async_streaming_chat_extracts_model_from_chunk(mock_client):
"""Test async streaming extracts model from chunk when not in kwargs."""
chunks = [
ChatCompletionChunk(
id="chunk1",
model="gpt-4o-async-stored",
object="chat.completion.chunk",
created=1234567890,
choices=[
ChoiceChunk(
index=0,
delta=ChoiceDelta(role="assistant", content="Hello"),
finish_reason="stop",
)
],
usage=CompletionUsage(
prompt_tokens=10,
completion_tokens=5,
total_tokens=15,
),
),
]
async def mock_create(self, **kwargs):
async def chunk_iterable():
for chunk in chunks:
yield chunk
return chunk_iterable()
with patch(
"openai.resources.chat.completions.AsyncCompletions.create", new=mock_create
):
client = AsyncOpenAI(api_key="test-key", insights_client=mock_client)
# Note: NOT passing model
response_stream = await client.chat.completions.create(
messages=[{"role": "user", "content": "Hello"}],
stream=True,
insights_distinct_id="test-id",
)
async for _ in response_stream:
pass
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
assert props["$ai_model"] == "gpt-4o-async-stored"
@pytest.mark.asyncio
async def test_async_streaming_responses_extracts_model_from_response(mock_client):
"""Test async Responses API streaming extracts model from chunk.response.model."""
from unittest.mock import MagicMock
from openai.types.responses import ResponseUsage
chunks = []
chunk1 = MagicMock()
chunk1.type = "response.text.delta"
chunk1.text = "Test"
del chunk1.response
chunks.append(chunk1)
chunk2 = MagicMock()
chunk2.type = "response.completed"
chunk2.response = MagicMock()
chunk2.response.model = "gpt-4o-mini-async-stored"
chunk2.response.usage = ResponseUsage(
input_tokens=20,
output_tokens=10,
total_tokens=30,
input_tokens_details={"prompt_tokens": 20, "cached_tokens": 0},
output_tokens_details={"reasoning_tokens": 0},
)
chunk2.response.output = ["Test"]
chunks.append(chunk2)
async def mock_create(self, **kwargs):
async def chunk_iterable():
for chunk in chunks:
yield chunk
return chunk_iterable()
with patch("openai.resources.responses.AsyncResponses.create", new=mock_create):
client = AsyncOpenAI(api_key="test-key", insights_client=mock_client)
response_stream = await client.responses.create(
input=[{"role": "user", "content": "Hello"}],
stream=True,
insights_distinct_id="test-id",
)
async for _ in response_stream:
pass
call_args = mock_client.capture.call_args[1]
props = call_args["properties"]
assert props["$ai_model"] == "gpt-4o-mini-async-stored"
@@ -0,0 +1 @@
# Tests for OpenAI Agents SDK integration
@@ -0,0 +1,810 @@
import logging
from unittest.mock import MagicMock, patch
import pytest
try:
from agents.tracing.span_data import (
AgentSpanData,
CustomSpanData,
FunctionSpanData,
GenerationSpanData,
GuardrailSpanData,
HandoffSpanData,
ResponseSpanData,
SpeechSpanData,
TranscriptionSpanData,
)
from hanzo_insights.ai.openai_agents import InsightsTracingProcessor, instrument
OPENAI_AGENTS_AVAILABLE = True
except ImportError:
OPENAI_AGENTS_AVAILABLE = False
# Skip all tests if OpenAI Agents SDK is not available
pytestmark = pytest.mark.skipif(
not OPENAI_AGENTS_AVAILABLE, reason="OpenAI Agents SDK is not available"
)
@pytest.fixture(scope="function")
def mock_client():
client = MagicMock()
client.privacy_mode = False
logging.getLogger("hanzo_insights").setLevel(logging.DEBUG)
return client
@pytest.fixture(scope="function")
def processor(mock_client):
return InsightsTracingProcessor(
client=mock_client,
distinct_id="test-user",
privacy_mode=False,
)
@pytest.fixture
def mock_trace():
trace = MagicMock()
trace.trace_id = "trace_123456789"
trace.name = "Test Workflow"
trace.group_id = "group_123"
trace.metadata = {"key": "value"}
return trace
@pytest.fixture
def mock_span():
span = MagicMock()
span.trace_id = "trace_123456789"
span.span_id = "span_987654321"
span.parent_id = None
span.started_at = "2024-01-01T00:00:00Z"
span.ended_at = "2024-01-01T00:00:01Z"
span.error = None
return span
class TestInsightsTracingProcessor:
"""Tests for the InsightsTracingProcessor class."""
def test_initialization(self, mock_client):
"""Test processor initializes correctly."""
processor = InsightsTracingProcessor(
client=mock_client,
distinct_id="user@example.com",
privacy_mode=True,
groups={"company": "acme"},
properties={"env": "test"},
)
assert processor._client == mock_client
assert processor._distinct_id == "user@example.com"
assert processor._privacy_mode is True
assert processor._groups == {"company": "acme"}
assert processor._properties == {"env": "test"}
def test_initialization_with_callable_distinct_id(self, mock_client, mock_trace):
"""Test processor with callable distinct_id resolver."""
def resolver(trace):
return trace.metadata.get("user_id", "default")
processor = InsightsTracingProcessor(
client=mock_client,
distinct_id=resolver,
)
mock_trace.metadata = {"user_id": "resolved-user"}
distinct_id = processor._get_distinct_id(mock_trace)
assert distinct_id == "resolved-user"
def test_on_trace_start_stores_metadata(self, processor, mock_client, mock_trace):
"""Test that on_trace_start stores metadata but does not capture an event."""
processor.on_trace_start(mock_trace)
mock_client.capture.assert_not_called()
assert mock_trace.trace_id in processor._trace_metadata
def test_on_trace_end_captures_ai_trace(self, processor, mock_client, mock_trace):
"""Test that on_trace_end captures $ai_trace event."""
processor.on_trace_start(mock_trace)
processor.on_trace_end(mock_trace)
mock_client.capture.assert_called_once()
call_kwargs = mock_client.capture.call_args[1]
assert call_kwargs["event"] == "$ai_trace"
assert call_kwargs["distinct_id"] == "test-user"
assert call_kwargs["properties"]["$ai_trace_id"] == "trace_123456789"
assert call_kwargs["properties"]["$ai_trace_name"] == "Test Workflow"
assert call_kwargs["properties"]["$ai_provider"] == "openai"
assert call_kwargs["properties"]["$ai_framework"] == "openai-agents"
assert "$ai_latency" in call_kwargs["properties"]
def test_personless_mode_when_no_distinct_id(self, mock_client, mock_trace):
"""Test that trace events use personless mode when no distinct_id is provided."""
processor = InsightsTracingProcessor(
client=mock_client,
)
processor.on_trace_start(mock_trace)
processor.on_trace_end(mock_trace)
call_kwargs = mock_client.capture.call_args[1]
assert call_kwargs["properties"]["$process_person_profile"] is False
# Should fallback to trace_id as the distinct_id
assert call_kwargs["distinct_id"] == mock_trace.trace_id
def test_personless_mode_for_spans_when_no_distinct_id(
self, mock_client, mock_trace, mock_span
):
"""Test that span events use personless mode when no distinct_id is provided."""
processor = InsightsTracingProcessor(
client=mock_client,
)
processor.on_trace_start(mock_trace)
mock_client.capture.reset_mock()
span_data = GenerationSpanData(model="gpt-4o")
mock_span.span_data = span_data
processor.on_span_start(mock_span)
processor.on_span_end(mock_span)
call_kwargs = mock_client.capture.call_args[1]
assert call_kwargs["properties"]["$process_person_profile"] is False
assert call_kwargs["distinct_id"] == mock_span.trace_id
def test_personless_mode_when_callable_returns_none(
self, mock_client, mock_trace, mock_span
):
"""Test personless mode when callable distinct_id returns None."""
def resolver(trace):
return None # Simulate no user ID available
processor = InsightsTracingProcessor(
client=mock_client,
distinct_id=resolver,
)
processor.on_trace_start(mock_trace)
mock_client.capture.reset_mock()
span_data = GenerationSpanData(model="gpt-4o")
mock_span.span_data = span_data
processor.on_span_start(mock_span)
processor.on_span_end(mock_span)
call_kwargs = mock_client.capture.call_args[1]
assert call_kwargs["properties"]["$process_person_profile"] is False
assert call_kwargs["distinct_id"] == mock_span.trace_id
def test_person_profile_when_distinct_id_provided(self, mock_client, mock_trace):
"""Test that events create person profiles when distinct_id is provided."""
processor = InsightsTracingProcessor(
client=mock_client,
distinct_id="real-user",
)
processor.on_trace_start(mock_trace)
processor.on_trace_end(mock_trace)
call_kwargs = mock_client.capture.call_args[1]
assert "$process_person_profile" not in call_kwargs["properties"]
def test_on_trace_end_clears_metadata(self, processor, mock_client, mock_trace):
"""Test that on_trace_end clears stored trace metadata."""
processor.on_trace_start(mock_trace)
assert mock_trace.trace_id in processor._trace_metadata
processor.on_trace_end(mock_trace)
assert mock_trace.trace_id not in processor._trace_metadata
# Also verify it captured the event
mock_client.capture.assert_called_once()
def test_on_span_start_tracks_time(self, processor, mock_span):
"""Test that on_span_start records start time."""
processor.on_span_start(mock_span)
assert mock_span.span_id in processor._span_start_times
def test_generation_span_mapping(self, processor, mock_client, mock_span):
"""Test GenerationSpanData maps to $ai_generation event."""
span_data = GenerationSpanData(
input=[{"role": "user", "content": "Hello"}],
output=[{"role": "assistant", "content": "Hi there!"}],
model="gpt-4o",
model_config={"temperature": 0.7, "max_tokens": 100},
usage={"input_tokens": 10, "output_tokens": 20},
)
mock_span.span_data = span_data
processor.on_span_start(mock_span)
processor.on_span_end(mock_span)
mock_client.capture.assert_called_once()
call_kwargs = mock_client.capture.call_args[1]
assert call_kwargs["event"] == "$ai_generation"
assert call_kwargs["properties"]["$ai_trace_id"] == "trace_123456789"
assert call_kwargs["properties"]["$ai_span_id"] == "span_987654321"
assert call_kwargs["properties"]["$ai_provider"] == "openai"
assert call_kwargs["properties"]["$ai_framework"] == "openai-agents"
assert call_kwargs["properties"]["$ai_model"] == "gpt-4o"
assert call_kwargs["properties"]["$ai_input_tokens"] == 10
assert call_kwargs["properties"]["$ai_output_tokens"] == 20
assert call_kwargs["properties"]["$ai_input"] == [
{"role": "user", "content": "Hello"}
]
assert call_kwargs["properties"]["$ai_output_choices"] == [
{"role": "assistant", "content": "Hi there!"}
]
def test_generation_span_with_reasoning_tokens(
self, processor, mock_client, mock_span
):
"""Test GenerationSpanData includes reasoning tokens when present."""
span_data = GenerationSpanData(
model="o1-preview",
usage={
"input_tokens": 100,
"output_tokens": 500,
"reasoning_tokens": 400,
},
)
mock_span.span_data = span_data
processor.on_span_start(mock_span)
processor.on_span_end(mock_span)
call_kwargs = mock_client.capture.call_args[1]
assert call_kwargs["properties"]["$ai_reasoning_tokens"] == 400
def test_function_span_mapping(self, processor, mock_client, mock_span):
"""Test FunctionSpanData maps to $ai_span event with type=tool."""
span_data = FunctionSpanData(
name="get_weather",
input='{"city": "San Francisco"}',
output="Sunny, 72F",
)
mock_span.span_data = span_data
processor.on_span_start(mock_span)
processor.on_span_end(mock_span)
call_kwargs = mock_client.capture.call_args[1]
assert call_kwargs["event"] == "$ai_span"
assert call_kwargs["properties"]["$ai_span_name"] == "get_weather"
assert call_kwargs["properties"]["$ai_span_type"] == "tool"
assert (
call_kwargs["properties"]["$ai_input_state"] == '{"city": "San Francisco"}'
)
assert call_kwargs["properties"]["$ai_output_state"] == "Sunny, 72F"
def test_agent_span_mapping(self, processor, mock_client, mock_span):
"""Test AgentSpanData maps to $ai_span event with type=agent."""
span_data = AgentSpanData(
name="CustomerServiceAgent",
handoffs=["TechnicalAgent", "BillingAgent"],
tools=["search", "get_order"],
output_type="str",
)
mock_span.span_data = span_data
processor.on_span_start(mock_span)
processor.on_span_end(mock_span)
call_kwargs = mock_client.capture.call_args[1]
assert call_kwargs["event"] == "$ai_span"
assert call_kwargs["properties"]["$ai_span_name"] == "CustomerServiceAgent"
assert call_kwargs["properties"]["$ai_span_type"] == "agent"
assert call_kwargs["properties"]["$ai_agent_handoffs"] == [
"TechnicalAgent",
"BillingAgent",
]
assert call_kwargs["properties"]["$ai_agent_tools"] == ["search", "get_order"]
def test_handoff_span_mapping(self, processor, mock_client, mock_span):
"""Test HandoffSpanData maps to $ai_span event with type=handoff."""
span_data = HandoffSpanData(
from_agent="TriageAgent",
to_agent="TechnicalAgent",
)
mock_span.span_data = span_data
processor.on_span_start(mock_span)
processor.on_span_end(mock_span)
call_kwargs = mock_client.capture.call_args[1]
assert call_kwargs["event"] == "$ai_span"
assert call_kwargs["properties"]["$ai_span_type"] == "handoff"
assert call_kwargs["properties"]["$ai_handoff_from_agent"] == "TriageAgent"
assert call_kwargs["properties"]["$ai_handoff_to_agent"] == "TechnicalAgent"
assert (
call_kwargs["properties"]["$ai_span_name"]
== "TriageAgent -> TechnicalAgent"
)
def test_guardrail_span_mapping(self, processor, mock_client, mock_span):
"""Test GuardrailSpanData maps to $ai_span event with type=guardrail."""
span_data = GuardrailSpanData(
name="ContentFilter",
triggered=True,
)
mock_span.span_data = span_data
processor.on_span_start(mock_span)
processor.on_span_end(mock_span)
call_kwargs = mock_client.capture.call_args[1]
assert call_kwargs["event"] == "$ai_span"
assert call_kwargs["properties"]["$ai_span_name"] == "ContentFilter"
assert call_kwargs["properties"]["$ai_span_type"] == "guardrail"
assert call_kwargs["properties"]["$ai_guardrail_triggered"] is True
def test_custom_span_mapping(self, processor, mock_client, mock_span):
"""Test CustomSpanData maps to $ai_span event with type=custom."""
span_data = CustomSpanData(
name="database_query",
data={"query": "SELECT * FROM users", "rows": 100},
)
mock_span.span_data = span_data
processor.on_span_start(mock_span)
processor.on_span_end(mock_span)
call_kwargs = mock_client.capture.call_args[1]
assert call_kwargs["event"] == "$ai_span"
assert call_kwargs["properties"]["$ai_span_name"] == "database_query"
assert call_kwargs["properties"]["$ai_span_type"] == "custom"
assert call_kwargs["properties"]["$ai_custom_data"] == {
"query": "SELECT * FROM users",
"rows": 100,
}
def test_privacy_mode_redacts_content(self, mock_client, mock_span):
"""Test that privacy_mode redacts input/output content."""
processor = InsightsTracingProcessor(
client=mock_client,
distinct_id="test-user",
privacy_mode=True,
)
span_data = GenerationSpanData(
input=[{"role": "user", "content": "Secret message"}],
output=[{"role": "assistant", "content": "Secret response"}],
model="gpt-4o",
usage={"input_tokens": 10, "output_tokens": 20},
)
mock_span.span_data = span_data
processor.on_span_start(mock_span)
processor.on_span_end(mock_span)
call_kwargs = mock_client.capture.call_args[1]
# Content should be redacted
assert call_kwargs["properties"]["$ai_input"] is None
assert call_kwargs["properties"]["$ai_output_choices"] is None
# Token counts should still be present
assert call_kwargs["properties"]["$ai_input_tokens"] == 10
assert call_kwargs["properties"]["$ai_output_tokens"] == 20
def test_error_handling_in_span(self, processor, mock_client, mock_span):
"""Test that span errors are captured correctly."""
span_data = GenerationSpanData(model="gpt-4o")
mock_span.span_data = span_data
mock_span.error = {"message": "Rate limit exceeded", "data": {"code": 429}}
processor.on_span_start(mock_span)
processor.on_span_end(mock_span)
call_kwargs = mock_client.capture.call_args[1]
assert call_kwargs["properties"]["$ai_is_error"] is True
assert call_kwargs["properties"]["$ai_error"] == "Rate limit exceeded"
def test_generation_span_includes_total_tokens(
self, processor, mock_client, mock_span
):
"""Test that $ai_total_tokens is calculated and included."""
span_data = GenerationSpanData(
model="gpt-4o",
usage={"input_tokens": 100, "output_tokens": 50},
)
mock_span.span_data = span_data
processor.on_span_start(mock_span)
processor.on_span_end(mock_span)
call_kwargs = mock_client.capture.call_args[1]
assert call_kwargs["properties"]["$ai_total_tokens"] == 150
def test_error_type_categorization_model_behavior(
self, processor, mock_client, mock_span
):
"""Test that ModelBehaviorError is categorized correctly."""
span_data = GenerationSpanData(model="gpt-4o")
mock_span.span_data = span_data
mock_span.error = {
"message": "ModelBehaviorError: Invalid JSON output",
"type": "ModelBehaviorError",
}
processor.on_span_start(mock_span)
processor.on_span_end(mock_span)
call_kwargs = mock_client.capture.call_args[1]
assert call_kwargs["properties"]["$ai_error_type"] == "model_behavior_error"
def test_error_type_categorization_user_error(
self, processor, mock_client, mock_span
):
"""Test that UserError is categorized correctly."""
span_data = GenerationSpanData(model="gpt-4o")
mock_span.span_data = span_data
mock_span.error = {"message": "UserError: Tool failed", "type": "UserError"}
processor.on_span_start(mock_span)
processor.on_span_end(mock_span)
call_kwargs = mock_client.capture.call_args[1]
assert call_kwargs["properties"]["$ai_error_type"] == "user_error"
def test_error_type_categorization_input_guardrail(
self, processor, mock_client, mock_span
):
"""Test that InputGuardrailTripwireTriggered is categorized correctly."""
span_data = GenerationSpanData(model="gpt-4o")
mock_span.span_data = span_data
mock_span.error = {
"message": "InputGuardrailTripwireTriggered: Content blocked"
}
processor.on_span_start(mock_span)
processor.on_span_end(mock_span)
call_kwargs = mock_client.capture.call_args[1]
assert (
call_kwargs["properties"]["$ai_error_type"] == "input_guardrail_triggered"
)
def test_error_type_categorization_output_guardrail(
self, processor, mock_client, mock_span
):
"""Test that OutputGuardrailTripwireTriggered is categorized correctly."""
span_data = GenerationSpanData(model="gpt-4o")
mock_span.span_data = span_data
mock_span.error = {
"message": "OutputGuardrailTripwireTriggered: Response blocked"
}
processor.on_span_start(mock_span)
processor.on_span_end(mock_span)
call_kwargs = mock_client.capture.call_args[1]
assert (
call_kwargs["properties"]["$ai_error_type"] == "output_guardrail_triggered"
)
def test_error_type_categorization_max_turns(
self, processor, mock_client, mock_span
):
"""Test that MaxTurnsExceeded is categorized correctly."""
span_data = GenerationSpanData(model="gpt-4o")
mock_span.span_data = span_data
mock_span.error = {"message": "MaxTurnsExceeded: Agent exceeded maximum turns"}
processor.on_span_start(mock_span)
processor.on_span_end(mock_span)
call_kwargs = mock_client.capture.call_args[1]
assert call_kwargs["properties"]["$ai_error_type"] == "max_turns_exceeded"
def test_error_type_categorization_unknown(self, processor, mock_client, mock_span):
"""Test that unknown errors are categorized as unknown."""
span_data = GenerationSpanData(model="gpt-4o")
mock_span.span_data = span_data
mock_span.error = {"message": "Some random error occurred"}
processor.on_span_start(mock_span)
processor.on_span_end(mock_span)
call_kwargs = mock_client.capture.call_args[1]
assert call_kwargs["properties"]["$ai_error_type"] == "unknown"
def test_response_span_with_output_and_total_tokens(
self, processor, mock_client, mock_span
):
"""Test ResponseSpanData includes output choices and total tokens."""
# Create a mock response object
mock_response = MagicMock()
mock_response.id = "resp_123"
mock_response.model = "gpt-4o"
mock_response.output = [{"type": "message", "content": "Hello!"}]
mock_response.usage = MagicMock()
mock_response.usage.input_tokens = 25
mock_response.usage.output_tokens = 10
span_data = ResponseSpanData(
response=mock_response,
input="Hello, world!",
)
mock_span.span_data = span_data
processor.on_span_start(mock_span)
processor.on_span_end(mock_span)
call_kwargs = mock_client.capture.call_args[1]
assert call_kwargs["event"] == "$ai_generation"
assert call_kwargs["properties"]["$ai_total_tokens"] == 35
assert call_kwargs["properties"]["$ai_output_choices"] == [
{"type": "message", "content": "Hello!"}
]
assert call_kwargs["properties"]["$ai_response_id"] == "resp_123"
def test_speech_span_with_pass_through_properties(
self, processor, mock_client, mock_span
):
"""Test SpeechSpanData includes pass-through properties."""
span_data = SpeechSpanData(
input="Hello, how can I help you?",
output="base64_audio_data",
output_format="pcm",
model="tts-1",
model_config={"voice": "alloy", "speed": 1.0},
first_content_at="2024-01-01T00:00:00.500Z",
)
mock_span.span_data = span_data
processor.on_span_start(mock_span)
processor.on_span_end(mock_span)
call_kwargs = mock_client.capture.call_args[1]
assert call_kwargs["event"] == "$ai_span"
assert call_kwargs["properties"]["$ai_span_type"] == "speech"
assert call_kwargs["properties"]["$ai_model"] == "tts-1"
# Pass-through properties (no $ai_ prefix)
assert (
call_kwargs["properties"]["first_content_at"] == "2024-01-01T00:00:00.500Z"
)
assert call_kwargs["properties"]["audio_output_format"] == "pcm"
assert call_kwargs["properties"]["model_config"] == {
"voice": "alloy",
"speed": 1.0,
}
# Text input should be captured
assert call_kwargs["properties"]["$ai_input"] == "Hello, how can I help you?"
def test_transcription_span_with_pass_through_properties(
self, processor, mock_client, mock_span
):
"""Test TranscriptionSpanData includes pass-through properties."""
span_data = TranscriptionSpanData(
input="base64_audio_data",
input_format="pcm",
output="This is the transcribed text.",
model="whisper-1",
model_config={"language": "en"},
)
mock_span.span_data = span_data
processor.on_span_start(mock_span)
processor.on_span_end(mock_span)
call_kwargs = mock_client.capture.call_args[1]
assert call_kwargs["event"] == "$ai_span"
assert call_kwargs["properties"]["$ai_span_type"] == "transcription"
assert call_kwargs["properties"]["$ai_model"] == "whisper-1"
# Pass-through properties (no $ai_ prefix)
assert call_kwargs["properties"]["audio_input_format"] == "pcm"
assert call_kwargs["properties"]["model_config"] == {"language": "en"}
# Transcription output should be captured
assert (
call_kwargs["properties"]["$ai_output_state"]
== "This is the transcribed text."
)
def test_latency_calculation(self, processor, mock_client, mock_span):
"""Test that latency is calculated correctly."""
span_data = GenerationSpanData(model="gpt-4o")
mock_span.span_data = span_data
with patch("time.time") as mock_time:
mock_time.return_value = 1000.0
processor.on_span_start(mock_span)
mock_time.return_value = 1001.5 # 1.5 seconds later
processor.on_span_end(mock_span)
call_kwargs = mock_client.capture.call_args[1]
assert call_kwargs["properties"]["$ai_latency"] == pytest.approx(1.5, rel=0.01)
def test_groups_included_in_events(self, mock_client, mock_trace, mock_span):
"""Test that groups are included in captured events."""
processor = InsightsTracingProcessor(
client=mock_client,
distinct_id="test-user",
groups={"company": "acme", "team": "engineering"},
)
processor.on_trace_start(mock_trace)
processor.on_trace_end(mock_trace)
call_kwargs = mock_client.capture.call_args[1]
assert call_kwargs["groups"] == {"company": "acme", "team": "engineering"}
def test_additional_properties_included(self, mock_client, mock_trace):
"""Test that additional properties are included in events."""
processor = InsightsTracingProcessor(
client=mock_client,
distinct_id="test-user",
properties={"environment": "production", "version": "1.0"},
)
processor.on_trace_start(mock_trace)
processor.on_trace_end(mock_trace)
call_kwargs = mock_client.capture.call_args[1]
assert call_kwargs["properties"]["environment"] == "production"
assert call_kwargs["properties"]["version"] == "1.0"
def test_shutdown_clears_state(self, processor):
"""Test that shutdown clears internal state."""
processor._span_start_times["span_1"] = 1000.0
processor._trace_metadata["trace_1"] = {"name": "test"}
processor.shutdown()
assert len(processor._span_start_times) == 0
assert len(processor._trace_metadata) == 0
def test_force_flush_calls_client_flush(self, processor, mock_client):
"""Test that force_flush calls client.flush()."""
processor.force_flush()
mock_client.flush.assert_called_once()
def test_generation_span_with_no_usage(self, processor, mock_client, mock_span):
"""Test GenerationSpanData with no usage data defaults to zero tokens."""
span_data = GenerationSpanData(model="gpt-4o")
mock_span.span_data = span_data
processor.on_span_start(mock_span)
processor.on_span_end(mock_span)
call_kwargs = mock_client.capture.call_args[1]
assert call_kwargs["properties"]["$ai_input_tokens"] == 0
assert call_kwargs["properties"]["$ai_output_tokens"] == 0
assert call_kwargs["properties"]["$ai_total_tokens"] == 0
def test_generation_span_with_partial_usage(
self, processor, mock_client, mock_span
):
"""Test GenerationSpanData with only input_tokens present."""
span_data = GenerationSpanData(
model="gpt-4o",
usage={"input_tokens": 42},
)
mock_span.span_data = span_data
processor.on_span_start(mock_span)
processor.on_span_end(mock_span)
call_kwargs = mock_client.capture.call_args[1]
assert call_kwargs["properties"]["$ai_input_tokens"] == 42
assert call_kwargs["properties"]["$ai_output_tokens"] == 0
assert call_kwargs["properties"]["$ai_total_tokens"] == 42
def test_error_type_categorization_by_type_field_only(
self, processor, mock_client, mock_span
):
"""Test error categorization works when only the type field matches."""
span_data = GenerationSpanData(model="gpt-4o")
mock_span.span_data = span_data
mock_span.error = {
"message": "Something went wrong",
"type": "ModelBehaviorError",
}
processor.on_span_start(mock_span)
processor.on_span_end(mock_span)
call_kwargs = mock_client.capture.call_args[1]
assert call_kwargs["properties"]["$ai_error_type"] == "model_behavior_error"
def test_distinct_id_resolved_from_trace_for_spans(
self, mock_client, mock_trace, mock_span
):
"""Test that spans use the distinct_id resolved at trace start."""
def resolver(trace):
return f"user-{trace.name}"
processor = InsightsTracingProcessor(
client=mock_client,
distinct_id=resolver,
)
# Start trace - this resolves and stores distinct_id
processor.on_trace_start(mock_trace)
mock_client.capture.reset_mock()
# End a span - should use the stored distinct_id from trace
span_data = GenerationSpanData(model="gpt-4o")
mock_span.span_data = span_data
processor.on_span_start(mock_span)
processor.on_span_end(mock_span)
call_kwargs = mock_client.capture.call_args[1]
assert call_kwargs["distinct_id"] == "user-Test Workflow"
def test_eviction_of_stale_entries(self, mock_client):
"""Test that stale entries are evicted when max is exceeded."""
processor = InsightsTracingProcessor(
client=mock_client,
distinct_id="test-user",
)
processor._max_tracked_entries = 10
# Fill beyond max
for i in range(15):
processor._span_start_times[f"span_{i}"] = float(i)
processor._trace_metadata[f"trace_{i}"] = {"name": f"trace_{i}"}
processor._evict_stale_entries()
# Should have evicted half
assert len(processor._span_start_times) <= 10
assert len(processor._trace_metadata) <= 10
class TestInstrumentHelper:
"""Tests for the instrument() convenience function."""
def test_instrument_registers_processor(self, mock_client):
"""Test that instrument() registers a processor."""
with patch("agents.tracing.add_trace_processor") as mock_add:
processor = instrument(
client=mock_client,
distinct_id="test-user",
)
mock_add.assert_called_once_with(processor)
assert isinstance(processor, InsightsTracingProcessor)
def test_instrument_with_privacy_mode(self, mock_client):
"""Test instrument() respects privacy_mode."""
with patch("agents.tracing.add_trace_processor"):
processor = instrument(
client=mock_client,
privacy_mode=True,
)
assert processor._privacy_mode is True
def test_instrument_with_groups_and_properties(self, mock_client):
"""Test instrument() accepts groups and properties."""
with patch("agents.tracing.add_trace_processor"):
processor = instrument(
client=mock_client,
groups={"company": "acme"},
properties={"env": "test"},
)
assert processor._groups == {"company": "acme"}
assert processor._properties == {"env": "test"}
+764
View File
@@ -0,0 +1,764 @@
import unittest
from unittest.mock import MagicMock, patch
from hanzo_insights.ai.prompts import Prompts
class MockResponse:
"""Mock HTTP response for testing."""
def __init__(self, json_data=None, status_code=200, ok=True):
self._json_data = json_data
self.status_code = status_code
self.ok = ok
def json(self):
if self._json_data is None:
raise ValueError("No JSON data")
return self._json_data
class TestPrompts(unittest.TestCase):
"""Tests for the Prompts class."""
mock_prompt_response = {
"id": 1,
"name": "test-prompt",
"prompt": "Hello, {{name}}! You are a helpful assistant for {{company}}.",
"version": 1,
"created_by": "user@example.com",
"created_at": "2024-01-01T00:00:00Z",
"updated_at": "2024-01-01T00:00:00Z",
"deleted": False,
}
def create_mock_client(
self,
personal_api_key="phx_test_key",
project_api_key="phc_test_key",
host="https://us.insights.hanzo.ai",
):
"""Create a mock Insights client."""
mock = MagicMock()
mock.personal_api_key = personal_api_key
mock.api_key = project_api_key
mock.raw_host = host
return mock
class TestPromptsGet(TestPrompts):
"""Tests for the Prompts.get() method."""
@patch("hanzo_insights.ai.prompts._get_session")
def test_successfully_fetch_a_prompt(self, mock_get_session):
"""Should successfully fetch a prompt."""
mock_get = mock_get_session.return_value.get
mock_get.return_value = MockResponse(json_data=self.mock_prompt_response)
client = self.create_mock_client()
prompts = Prompts(client)
result = prompts.get("test-prompt")
self.assertEqual(result, self.mock_prompt_response["prompt"])
mock_get.assert_called_once()
call_args = mock_get.call_args
self.assertEqual(
call_args[0][0],
"https://us.insights.hanzo.ai/api/environments/@current/llm_prompts/name/test-prompt/?token=phc_test_key",
)
self.assertIn("Authorization", call_args[1]["headers"])
self.assertEqual(
call_args[1]["headers"]["Authorization"], "Bearer phx_test_key"
)
@patch("hanzo_insights.ai.prompts._get_session")
def test_successfully_fetch_a_specific_prompt_version(self, mock_get_session):
"""Should successfully fetch a specific prompt version."""
mock_get = mock_get_session.return_value.get
versioned_prompt_response = {
**self.mock_prompt_response,
"prompt": "Prompt version 1",
"version": 1,
}
mock_get.return_value = MockResponse(json_data=versioned_prompt_response)
client = self.create_mock_client()
prompts = Prompts(client)
result = prompts.get("test-prompt", version=1)
self.assertEqual(result, versioned_prompt_response["prompt"])
mock_get.assert_called_once()
call_args = mock_get.call_args
self.assertEqual(
call_args[0][0],
"https://us.insights.hanzo.ai/api/environments/@current/llm_prompts/name/test-prompt/?token=phc_test_key&version=1",
)
@patch("hanzo_insights.ai.prompts._get_session")
@patch("hanzo_insights.ai.prompts.time.time")
def test_return_cached_prompt_when_fresh(self, mock_time, mock_get_session):
"""Should return cached prompt when fresh (no API call)."""
mock_get = mock_get_session.return_value.get
mock_get.return_value = MockResponse(json_data=self.mock_prompt_response)
mock_time.return_value = 1000.0
client = self.create_mock_client()
prompts = Prompts(client)
# First call - fetches from API
result1 = prompts.get("test-prompt", cache_ttl_seconds=300)
self.assertEqual(result1, self.mock_prompt_response["prompt"])
self.assertEqual(mock_get.call_count, 1)
# Advance time by 60 seconds (still within TTL)
mock_time.return_value = 1060.0
# Second call - should use cache
result2 = prompts.get("test-prompt", cache_ttl_seconds=300)
self.assertEqual(result2, self.mock_prompt_response["prompt"])
self.assertEqual(mock_get.call_count, 1) # No additional fetch
@patch("hanzo_insights.ai.prompts._get_session")
def test_cache_latest_and_versioned_prompts_separately(self, mock_get_session):
"""Should cache latest and historical prompt versions separately."""
mock_get = mock_get_session.return_value.get
latest_prompt_response = {
**self.mock_prompt_response,
"prompt": "Latest prompt",
"version": 2,
}
versioned_prompt_response = {
**self.mock_prompt_response,
"prompt": "Prompt version 1",
"version": 1,
}
mock_get.side_effect = [
MockResponse(json_data=latest_prompt_response),
MockResponse(json_data=versioned_prompt_response),
]
client = self.create_mock_client()
prompts = Prompts(client)
self.assertEqual(prompts.get("test-prompt"), latest_prompt_response["prompt"])
self.assertEqual(
prompts.get("test-prompt", version=1),
versioned_prompt_response["prompt"],
)
self.assertEqual(prompts.get("test-prompt"), latest_prompt_response["prompt"])
self.assertEqual(
prompts.get("test-prompt", version=1),
versioned_prompt_response["prompt"],
)
self.assertEqual(mock_get.call_count, 2)
@patch("hanzo_insights.ai.prompts._get_session")
@patch("hanzo_insights.ai.prompts.time.time")
def test_refetch_when_cache_is_stale(self, mock_time, mock_get_session):
"""Should refetch when cache is stale."""
mock_get = mock_get_session.return_value.get
updated_prompt_response = {
**self.mock_prompt_response,
"prompt": "Updated prompt: Hello, {{name}}!",
}
mock_get.side_effect = [
MockResponse(json_data=self.mock_prompt_response),
MockResponse(json_data=updated_prompt_response),
]
mock_time.return_value = 1000.0
client = self.create_mock_client()
prompts = Prompts(client)
# First call - fetches from API
result1 = prompts.get("test-prompt", cache_ttl_seconds=60)
self.assertEqual(result1, self.mock_prompt_response["prompt"])
self.assertEqual(mock_get.call_count, 1)
# Advance time past TTL
mock_time.return_value = 1061.0
# Second call - should refetch
result2 = prompts.get("test-prompt", cache_ttl_seconds=60)
self.assertEqual(result2, updated_prompt_response["prompt"])
self.assertEqual(mock_get.call_count, 2)
@patch("hanzo_insights.ai.prompts._get_session")
@patch("hanzo_insights.ai.prompts.time.time")
@patch("hanzo_insights.ai.prompts.log")
def test_use_stale_cache_on_fetch_failure_with_warning(
self, mock_log, mock_time, mock_get_session
):
"""Should use stale cache on fetch failure with warning."""
mock_get = mock_get_session.return_value.get
mock_get.side_effect = [
MockResponse(json_data=self.mock_prompt_response),
Exception("Network error"),
]
mock_time.return_value = 1000.0
client = self.create_mock_client()
prompts = Prompts(client)
# First call - populates cache
result1 = prompts.get("test-prompt", cache_ttl_seconds=60)
self.assertEqual(result1, self.mock_prompt_response["prompt"])
# Advance time past TTL
mock_time.return_value = 1061.0
# Second call - should use stale cache
result2 = prompts.get("test-prompt", cache_ttl_seconds=60)
self.assertEqual(result2, self.mock_prompt_response["prompt"])
# Check warning was logged
mock_log.warning.assert_called()
warning_call = mock_log.warning.call_args
self.assertIn("using stale cache", warning_call[0][0])
@patch("hanzo_insights.ai.prompts._get_session")
@patch("hanzo_insights.ai.prompts.log")
def test_use_fallback_when_no_cache_and_fetch_fails_with_warning(
self, mock_log, mock_get_session
):
"""Should use fallback when no cache and fetch fails with warning."""
mock_get = mock_get_session.return_value.get
mock_get.side_effect = Exception("Network error")
client = self.create_mock_client()
prompts = Prompts(client)
fallback = "Default system prompt."
result = prompts.get("test-prompt", fallback=fallback)
self.assertEqual(result, fallback)
# Check warning was logged
mock_log.warning.assert_called()
warning_call = mock_log.warning.call_args
self.assertIn("using fallback", warning_call[0][0])
@patch("hanzo_insights.ai.prompts._get_session")
def test_throw_when_no_cache_no_fallback_and_fetch_fails(self, mock_get_session):
"""Should throw when no cache, no fallback, and fetch fails."""
mock_get = mock_get_session.return_value.get
mock_get.side_effect = Exception("Network error")
client = self.create_mock_client()
prompts = Prompts(client)
with self.assertRaises(Exception) as context:
prompts.get("test-prompt")
self.assertIn("Network error", str(context.exception))
@patch("hanzo_insights.ai.prompts._get_session")
def test_handle_404_response(self, mock_get_session):
"""Should handle 404 response."""
mock_get = mock_get_session.return_value.get
mock_get.return_value = MockResponse(status_code=404, ok=False)
client = self.create_mock_client()
prompts = Prompts(client)
with self.assertRaises(Exception) as context:
prompts.get("nonexistent-prompt")
self.assertIn('Prompt "nonexistent-prompt" not found', str(context.exception))
@patch("hanzo_insights.ai.prompts._get_session")
def test_handle_404_response_for_specific_prompt_version(self, mock_get_session):
"""Should handle 404 response for a specific prompt version."""
mock_get = mock_get_session.return_value.get
mock_get.return_value = MockResponse(status_code=404, ok=False)
client = self.create_mock_client()
prompts = Prompts(client)
with self.assertRaises(Exception) as context:
prompts.get("nonexistent-prompt", version=3)
self.assertIn(
'Prompt "nonexistent-prompt" version 3 not found',
str(context.exception),
)
@patch("hanzo_insights.ai.prompts._get_session")
def test_handle_403_response(self, mock_get_session):
"""Should handle 403 response."""
mock_get = mock_get_session.return_value.get
mock_get.return_value = MockResponse(status_code=403, ok=False)
client = self.create_mock_client()
prompts = Prompts(client)
with self.assertRaises(Exception) as context:
prompts.get("restricted-prompt")
self.assertIn(
'Access denied for prompt "restricted-prompt"', str(context.exception)
)
def test_throw_when_no_personal_api_key_configured(self):
"""Should throw when no personal_api_key is configured."""
client = self.create_mock_client(personal_api_key=None)
prompts = Prompts(client)
with self.assertRaises(Exception) as context:
prompts.get("test-prompt")
self.assertIn(
"personal_api_key is required to fetch prompts", str(context.exception)
)
def test_throw_when_no_project_api_key_configured(self):
"""Should throw when no project_api_key is configured."""
client = self.create_mock_client(project_api_key=None)
prompts = Prompts(client)
with self.assertRaises(Exception) as context:
prompts.get("test-prompt")
self.assertIn(
"project_api_key is required to fetch prompts", str(context.exception)
)
@patch("hanzo_insights.ai.prompts._get_session")
def test_throw_when_api_returns_invalid_response_format(self, mock_get_session):
"""Should throw when API returns invalid response format."""
mock_get = mock_get_session.return_value.get
mock_get.return_value = MockResponse(json_data={"invalid": "response"})
client = self.create_mock_client()
prompts = Prompts(client)
with self.assertRaises(Exception) as context:
prompts.get("test-prompt")
self.assertIn("Invalid response format", str(context.exception))
@patch("hanzo_insights.ai.prompts._get_session")
def test_use_custom_host_from_insights_options(self, mock_get_session):
"""Should use custom host from Insights options."""
mock_get = mock_get_session.return_value.get
mock_get.return_value = MockResponse(json_data=self.mock_prompt_response)
client = self.create_mock_client(host="https://eu.insights.hanzo.ai")
prompts = Prompts(client)
prompts.get("test-prompt")
call_args = mock_get.call_args
self.assertTrue(
call_args[0][0].startswith(
"https://eu.insights.hanzo.ai/api/environments/@current/llm_prompts/name/test-prompt/?token=phc_test_key"
),
f"Expected URL to start with 'https://eu.insights.hanzo.ai/api/environments/@current/llm_prompts/name/test-prompt/?token=phc_test_key', got {call_args[0][0]}",
)
@patch("hanzo_insights.ai.prompts._get_session")
@patch("hanzo_insights.ai.prompts.time.time")
def test_use_default_cache_ttl_5_minutes(self, mock_time, mock_get_session):
"""Should use default cache TTL (5 minutes) when not specified."""
mock_get = mock_get_session.return_value.get
mock_get.return_value = MockResponse(json_data=self.mock_prompt_response)
mock_time.return_value = 1000.0
client = self.create_mock_client()
prompts = Prompts(client)
# First call
prompts.get("test-prompt")
self.assertEqual(mock_get.call_count, 1)
# Advance time by 4 minutes (within default 5-minute TTL)
mock_time.return_value = 1000.0 + (4 * 60)
# Second call - should use cache
prompts.get("test-prompt")
self.assertEqual(mock_get.call_count, 1)
# Advance time past 5-minute TTL
mock_time.return_value = 1000.0 + (6 * 60)
# Third call - should refetch
prompts.get("test-prompt")
self.assertEqual(mock_get.call_count, 2)
@patch("hanzo_insights.ai.prompts._get_session")
@patch("hanzo_insights.ai.prompts.time.time")
def test_use_custom_default_cache_ttl_from_constructor(
self, mock_time, mock_get_session
):
"""Should use custom default cache TTL from constructor."""
mock_get = mock_get_session.return_value.get
mock_get.return_value = MockResponse(json_data=self.mock_prompt_response)
mock_time.return_value = 1000.0
client = self.create_mock_client()
prompts = Prompts(client, default_cache_ttl_seconds=60)
# First call
prompts.get("test-prompt")
self.assertEqual(mock_get.call_count, 1)
# Advance time past custom TTL
mock_time.return_value = 1061.0
# Second call - should refetch
prompts.get("test-prompt")
self.assertEqual(mock_get.call_count, 2)
@patch("hanzo_insights.ai.prompts._get_session")
def test_url_encode_prompt_names_with_special_characters(self, mock_get_session):
"""Should URL-encode prompt names with special characters."""
mock_get = mock_get_session.return_value.get
mock_get.return_value = MockResponse(json_data=self.mock_prompt_response)
client = self.create_mock_client()
prompts = Prompts(client)
prompts.get("prompt with spaces/and/slashes")
call_args = mock_get.call_args
self.assertEqual(
call_args[0][0],
"https://us.insights.hanzo.ai/api/environments/@current/llm_prompts/name/prompt%20with%20spaces%2Fand%2Fslashes/?token=phc_test_key",
)
@patch("hanzo_insights.ai.prompts._get_session")
def test_work_with_direct_options_no_insights_client(self, mock_get_session):
"""Should work with direct options (no Insights client)."""
mock_get = mock_get_session.return_value.get
mock_get.return_value = MockResponse(json_data=self.mock_prompt_response)
prompts = Prompts(
personal_api_key="phx_direct_key", project_api_key="phc_direct_key"
)
result = prompts.get("test-prompt")
self.assertEqual(result, self.mock_prompt_response["prompt"])
call_args = mock_get.call_args
self.assertEqual(
call_args[0][0],
"https://us.insights.hanzo.ai/api/environments/@current/llm_prompts/name/test-prompt/?token=phc_direct_key",
)
self.assertEqual(
call_args[1]["headers"]["Authorization"], "Bearer phx_direct_key"
)
@patch("hanzo_insights.ai.prompts._get_session")
def test_use_custom_host_from_direct_options(self, mock_get_session):
"""Should use custom host from direct options."""
mock_get = mock_get_session.return_value.get
mock_get.return_value = MockResponse(json_data=self.mock_prompt_response)
prompts = Prompts(
personal_api_key="phx_direct_key",
project_api_key="phc_direct_key",
host="https://eu.insights.hanzo.ai",
)
prompts.get("test-prompt")
call_args = mock_get.call_args
self.assertEqual(
call_args[0][0],
"https://eu.insights.hanzo.ai/api/environments/@current/llm_prompts/name/test-prompt/?token=phc_direct_key",
)
@patch("hanzo_insights.ai.prompts._get_session")
@patch("hanzo_insights.ai.prompts.time.time")
def test_use_custom_default_cache_ttl_from_direct_options(
self, mock_time, mock_get_session
):
"""Should use custom default cache TTL from direct options."""
mock_get = mock_get_session.return_value.get
mock_get.return_value = MockResponse(json_data=self.mock_prompt_response)
mock_time.return_value = 1000.0
prompts = Prompts(
personal_api_key="phx_direct_key",
project_api_key="phc_direct_key",
default_cache_ttl_seconds=60,
)
# First call
prompts.get("test-prompt")
self.assertEqual(mock_get.call_count, 1)
# Advance time past custom TTL
mock_time.return_value = 1061.0
# Second call - should refetch
prompts.get("test-prompt")
self.assertEqual(mock_get.call_count, 2)
class TestPromptsCompile(TestPrompts):
"""Tests for the Prompts.compile() method."""
def test_replace_a_single_variable(self):
"""Should replace a single variable."""
client = self.create_mock_client()
prompts = Prompts(client)
result = prompts.compile("Hello, {{name}}!", {"name": "World"})
self.assertEqual(result, "Hello, World!")
def test_replace_multiple_variables(self):
"""Should replace multiple variables."""
client = self.create_mock_client()
prompts = Prompts(client)
result = prompts.compile(
"Hello, {{name}}! Welcome to {{company}}. Your tier is {{tier}}.",
{"name": "John", "company": "Acme Corp", "tier": "premium"},
)
self.assertEqual(
result, "Hello, John! Welcome to Acme Corp. Your tier is premium."
)
def test_handle_numbers(self):
"""Should handle numbers."""
client = self.create_mock_client()
prompts = Prompts(client)
result = prompts.compile("You have {{count}} items.", {"count": 42})
self.assertEqual(result, "You have 42 items.")
def test_handle_booleans(self):
"""Should handle booleans."""
client = self.create_mock_client()
prompts = Prompts(client)
result = prompts.compile("Feature enabled: {{enabled}}", {"enabled": True})
self.assertEqual(result, "Feature enabled: True")
def test_leave_unmatched_variables_unchanged(self):
"""Should leave unmatched variables unchanged."""
client = self.create_mock_client()
prompts = Prompts(client)
result = prompts.compile(
"Hello, {{name}}! Your {{unknown}} is ready.", {"name": "World"}
)
self.assertEqual(result, "Hello, World! Your {{unknown}} is ready.")
def test_handle_prompts_with_no_variables(self):
"""Should handle prompts with no variables."""
client = self.create_mock_client()
prompts = Prompts(client)
result = prompts.compile("You are a helpful assistant.", {})
self.assertEqual(result, "You are a helpful assistant.")
def test_handle_empty_variables_dict(self):
"""Should handle empty variables dict."""
client = self.create_mock_client()
prompts = Prompts(client)
result = prompts.compile("Hello, {{name}}!", {})
self.assertEqual(result, "Hello, {{name}}!")
def test_handle_multiple_occurrences_of_same_variable(self):
"""Should handle multiple occurrences of the same variable."""
client = self.create_mock_client()
prompts = Prompts(client)
result = prompts.compile(
"Hello, {{name}}! Goodbye, {{name}}!", {"name": "World"}
)
self.assertEqual(result, "Hello, World! Goodbye, World!")
def test_work_with_direct_options_initialization(self):
"""Should work with direct options initialization."""
prompts = Prompts(
personal_api_key="phx_test_key", project_api_key="phc_test_key"
)
result = prompts.compile("Hello, {{name}}!", {"name": "World"})
self.assertEqual(result, "Hello, World!")
def test_handle_variables_with_hyphens(self):
"""Should handle variables with hyphens."""
prompts = Prompts(
personal_api_key="phx_test_key", project_api_key="phc_test_key"
)
result = prompts.compile("User ID: {{user-id}}", {"user-id": "12345"})
self.assertEqual(result, "User ID: 12345")
def test_handle_variables_with_dots(self):
"""Should handle variables with dots."""
prompts = Prompts(
personal_api_key="phx_test_key", project_api_key="phc_test_key"
)
result = prompts.compile("Company: {{company.name}}", {"company.name": "Acme"})
self.assertEqual(result, "Company: Acme")
class TestPromptsClearCache(TestPrompts):
"""Tests for the Prompts.clear_cache() method."""
def test_clear_cache_with_version_and_no_name_raises_value_error(self):
"""Should enforce that versioned cache clearing requires a prompt name."""
client = self.create_mock_client()
prompts = Prompts(client)
with self.assertRaises(ValueError) as context:
prompts.clear_cache(version=1)
self.assertIn("requires 'name'", str(context.exception))
@patch("hanzo_insights.ai.prompts._get_session")
def test_clear_a_specific_prompt_from_cache(self, mock_get_session):
"""Should clear a specific prompt from cache."""
mock_get = mock_get_session.return_value.get
other_prompt_response = {**self.mock_prompt_response, "name": "other-prompt"}
mock_get.side_effect = [
MockResponse(json_data=self.mock_prompt_response),
MockResponse(json_data=other_prompt_response),
MockResponse(json_data=self.mock_prompt_response),
]
client = self.create_mock_client()
prompts = Prompts(client)
# Populate cache with two prompts
prompts.get("test-prompt")
prompts.get("other-prompt")
self.assertEqual(mock_get.call_count, 2)
# Clear only test-prompt
prompts.clear_cache("test-prompt")
# test-prompt should be refetched
prompts.get("test-prompt")
self.assertEqual(mock_get.call_count, 3)
# other-prompt should still be cached
prompts.get("other-prompt")
self.assertEqual(mock_get.call_count, 3)
@patch("hanzo_insights.ai.prompts._get_session")
def test_clear_a_specific_prompt_version_from_cache(self, mock_get_session):
"""Should clear only the requested prompt version from cache."""
mock_get = mock_get_session.return_value.get
latest_prompt_response = {
**self.mock_prompt_response,
"prompt": "Latest prompt",
"version": 2,
}
versioned_prompt_response = {
**self.mock_prompt_response,
"prompt": "Prompt version 1",
"version": 1,
}
mock_get.side_effect = [
MockResponse(json_data=latest_prompt_response),
MockResponse(json_data=versioned_prompt_response),
MockResponse(json_data=versioned_prompt_response),
]
client = self.create_mock_client()
prompts = Prompts(client)
prompts.get("test-prompt")
prompts.get("test-prompt", version=1)
self.assertEqual(mock_get.call_count, 2)
prompts.clear_cache("test-prompt", version=1)
prompts.get("test-prompt")
self.assertEqual(mock_get.call_count, 2)
prompts.get("test-prompt", version=1)
self.assertEqual(mock_get.call_count, 3)
@patch("hanzo_insights.ai.prompts._get_session")
def test_clear_a_prompt_name_clears_all_cached_versions(self, mock_get_session):
"""Should clear latest and versioned cache entries for the same prompt name."""
mock_get = mock_get_session.return_value.get
latest_prompt_response = {
**self.mock_prompt_response,
"prompt": "Latest prompt",
"version": 2,
}
versioned_prompt_response = {
**self.mock_prompt_response,
"prompt": "Prompt version 1",
"version": 1,
}
mock_get.side_effect = [
MockResponse(json_data=latest_prompt_response),
MockResponse(json_data=versioned_prompt_response),
MockResponse(json_data=latest_prompt_response),
MockResponse(json_data=versioned_prompt_response),
]
client = self.create_mock_client()
prompts = Prompts(client)
prompts.get("test-prompt")
prompts.get("test-prompt", version=1)
self.assertEqual(mock_get.call_count, 2)
prompts.clear_cache("test-prompt")
prompts.get("test-prompt")
prompts.get("test-prompt", version=1)
self.assertEqual(mock_get.call_count, 4)
@patch("hanzo_insights.ai.prompts._get_session")
def test_clear_all_prompts_from_cache(self, mock_get_session):
"""Should clear all prompts from cache when no name is provided."""
mock_get = mock_get_session.return_value.get
other_prompt_response = {**self.mock_prompt_response, "name": "other-prompt"}
mock_get.side_effect = [
MockResponse(json_data=self.mock_prompt_response),
MockResponse(json_data=other_prompt_response),
MockResponse(json_data=self.mock_prompt_response),
MockResponse(json_data=other_prompt_response),
]
client = self.create_mock_client()
prompts = Prompts(client)
# Populate cache with two prompts
prompts.get("test-prompt")
prompts.get("other-prompt")
self.assertEqual(mock_get.call_count, 2)
# Clear all cache
prompts.clear_cache()
# Both prompts should be refetched
prompts.get("test-prompt")
prompts.get("other-prompt")
self.assertEqual(mock_get.call_count, 4)
if __name__ == "__main__":
unittest.main()
@@ -1,6 +1,7 @@
import os
import unittest
from posthog.ai.sanitization import (
from hanzo_insights.ai.sanitization import (
redact_base64_data_url,
sanitize_openai,
sanitize_openai_response,
@@ -68,6 +69,25 @@ class TestSanitization(unittest.TestCase):
)
self.assertEqual(result[0]["content"][1]["image_url"]["detail"], "high")
def test_sanitize_openai_input_image(self):
input_data = [
{
"role": "user",
"content": [
{
"type": "input_image",
"image_url": self.sample_base64_image,
}
],
}
]
result = sanitize_openai(input_data)
self.assertEqual(
result[0]["content"][0]["image_url"], REDACTED_IMAGE_PLACEHOLDER
)
def test_sanitize_openai_preserves_regular_urls(self):
input_data = [
{
@@ -331,5 +351,191 @@ class TestSanitization(unittest.TestCase):
)
class TestAIMultipartRequest(unittest.TestCase):
"""Test that _INTERNAL_LLMA_MULTIMODAL environment variable controls sanitization."""
def tearDown(self):
# Clean up environment variable after each test
if "_INTERNAL_LLMA_MULTIMODAL" in os.environ:
del os.environ["_INTERNAL_LLMA_MULTIMODAL"]
def test_multimodal_disabled_redacts_images(self):
"""When _INTERNAL_LLMA_MULTIMODAL is not set, images should be redacted."""
if "_INTERNAL_LLMA_MULTIMODAL" in os.environ:
del os.environ["_INTERNAL_LLMA_MULTIMODAL"]
base64_image = "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQ..."
result = redact_base64_data_url(base64_image)
self.assertEqual(result, REDACTED_IMAGE_PLACEHOLDER)
def test_multimodal_enabled_preserves_images(self):
"""When _INTERNAL_LLMA_MULTIMODAL is true, images should be preserved."""
os.environ["_INTERNAL_LLMA_MULTIMODAL"] = "true"
base64_image = "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQ..."
result = redact_base64_data_url(base64_image)
self.assertEqual(result, base64_image)
def test_multimodal_enabled_with_1(self):
"""_INTERNAL_LLMA_MULTIMODAL=1 should enable multimodal."""
os.environ["_INTERNAL_LLMA_MULTIMODAL"] = "1"
base64_image = "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQ..."
result = redact_base64_data_url(base64_image)
self.assertEqual(result, base64_image)
def test_multimodal_enabled_with_yes(self):
"""_INTERNAL_LLMA_MULTIMODAL=yes should enable multimodal."""
os.environ["_INTERNAL_LLMA_MULTIMODAL"] = "yes"
base64_image = "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQ..."
result = redact_base64_data_url(base64_image)
self.assertEqual(result, base64_image)
def test_multimodal_false_redacts_images(self):
"""_INTERNAL_LLMA_MULTIMODAL=false should still redact."""
os.environ["_INTERNAL_LLMA_MULTIMODAL"] = "false"
base64_image = "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQ..."
result = redact_base64_data_url(base64_image)
self.assertEqual(result, REDACTED_IMAGE_PLACEHOLDER)
def test_anthropic_multimodal_enabled(self):
"""Anthropic images should be preserved when _INTERNAL_LLMA_MULTIMODAL is enabled."""
os.environ["_INTERNAL_LLMA_MULTIMODAL"] = "true"
input_data = [
{
"role": "user",
"content": [
{
"type": "image",
"source": {
"type": "base64",
"media_type": "image/jpeg",
"data": "base64data",
},
}
],
}
]
result = sanitize_anthropic(input_data)
self.assertEqual(result[0]["content"][0]["source"]["data"], "base64data")
def test_gemini_multimodal_enabled(self):
"""Gemini images should be preserved when _INTERNAL_LLMA_MULTIMODAL is enabled."""
os.environ["_INTERNAL_LLMA_MULTIMODAL"] = "true"
input_data = [
{
"parts": [
{"inline_data": {"mime_type": "image/jpeg", "data": "base64data"}}
]
}
]
result = sanitize_gemini(input_data)
self.assertEqual(result[0]["parts"][0]["inline_data"]["data"], "base64data")
def test_langchain_anthropic_style_multimodal_enabled(self):
"""LangChain Anthropic-style images should be preserved when _INTERNAL_LLMA_MULTIMODAL is enabled."""
os.environ["_INTERNAL_LLMA_MULTIMODAL"] = "true"
input_data = [
{
"role": "user",
"content": [
{
"type": "image",
"source": {"data": "base64data"},
}
],
}
]
result = sanitize_langchain(input_data)
self.assertEqual(result[0]["content"][0]["source"]["data"], "base64data")
def test_openai_audio_redacted_by_default(self):
"""OpenAI audio should be redacted when _INTERNAL_LLMA_MULTIMODAL is not set."""
if "_INTERNAL_LLMA_MULTIMODAL" in os.environ:
del os.environ["_INTERNAL_LLMA_MULTIMODAL"]
input_data = [
{
"role": "assistant",
"content": [
{"type": "audio", "data": "base64audiodata", "id": "audio_123"}
],
}
]
result = sanitize_openai(input_data)
self.assertEqual(result[0]["content"][0]["data"], REDACTED_IMAGE_PLACEHOLDER)
self.assertEqual(result[0]["content"][0]["id"], "audio_123")
def test_openai_audio_preserved_with_flag(self):
"""OpenAI audio should be preserved when _INTERNAL_LLMA_MULTIMODAL is enabled."""
os.environ["_INTERNAL_LLMA_MULTIMODAL"] = "true"
input_data = [
{
"role": "assistant",
"content": [
{"type": "audio", "data": "base64audiodata", "id": "audio_123"}
],
}
]
result = sanitize_openai(input_data)
self.assertEqual(result[0]["content"][0]["data"], "base64audiodata")
def test_gemini_audio_redacted_by_default(self):
"""Gemini audio should be redacted when _INTERNAL_LLMA_MULTIMODAL is not set."""
if "_INTERNAL_LLMA_MULTIMODAL" in os.environ:
del os.environ["_INTERNAL_LLMA_MULTIMODAL"]
input_data = [
{
"parts": [
{
"inline_data": {
"mime_type": "audio/L16;codec=pcm;rate=24000",
"data": "base64audiodata",
}
}
]
}
]
result = sanitize_gemini(input_data)
self.assertEqual(
result[0]["parts"][0]["inline_data"]["data"], REDACTED_IMAGE_PLACEHOLDER
)
def test_gemini_audio_preserved_with_flag(self):
"""Gemini audio should be preserved when _INTERNAL_LLMA_MULTIMODAL is enabled."""
os.environ["_INTERNAL_LLMA_MULTIMODAL"] = "true"
input_data = [
{
"parts": [
{
"inline_data": {
"mime_type": "audio/L16;codec=pcm;rate=24000",
"data": "base64audiodata",
}
}
]
}
]
result = sanitize_gemini(input_data)
self.assertEqual(
result[0]["parts"][0]["inline_data"]["data"], "base64audiodata"
)
if __name__ == "__main__":
unittest.main()
@@ -11,7 +11,10 @@ regardless of how they're passed to the providers:
import time
import unittest
from unittest.mock import patch, MagicMock
from unittest.mock import MagicMock, patch
from hanzo_insights.client import Client
from hanzo_insights.test.test_utils import FAKE_TEST_API_KEY
class TestSystemPromptCapture(unittest.TestCase):
@@ -23,8 +26,9 @@ class TestSystemPromptCapture(unittest.TestCase):
self.test_user_message = "Hello, how are you?"
self.test_response = "I'm doing well, thank you!"
# Create mock PostHog client
self.client = MagicMock()
# Create mock Insights client
self.client = Client(FAKE_TEST_API_KEY)
self.client._enqueue = MagicMock()
self.client.privacy_mode = False
def _assert_system_prompt_captured(self, captured_input):
@@ -53,10 +57,11 @@ class TestSystemPromptCapture(unittest.TestCase):
def test_openai_messages_array_system_prompt(self):
"""Test OpenAI with system prompt in messages array."""
try:
from posthog.ai.openai import OpenAI
from openai.types.chat import ChatCompletion, ChatCompletionMessage
from openai.types.chat.chat_completion import Choice
from openai.types.completion_usage import CompletionUsage
from hanzo_insights.ai.openai import OpenAI
except ImportError:
self.skipTest("OpenAI package not available")
@@ -83,7 +88,7 @@ class TestSystemPromptCapture(unittest.TestCase):
"openai.resources.chat.completions.Completions.create",
return_value=mock_response,
):
client = OpenAI(posthog_client=self.client, api_key="test")
client = OpenAI(insights_client=self.client, api_key="test")
messages = [
{"role": "system", "content": self.test_system_prompt},
@@ -91,20 +96,21 @@ class TestSystemPromptCapture(unittest.TestCase):
]
client.chat.completions.create(
model="gpt-4", messages=messages, posthog_distinct_id="test-user"
model="gpt-4", messages=messages, insights_distinct_id="test-user"
)
self.assertEqual(len(self.client.capture.call_args_list), 1)
properties = self.client.capture.call_args_list[0][1]["properties"]
self.assertEqual(len(self.client._enqueue.call_args_list), 1)
properties = self.client._enqueue.call_args_list[0][0][0]["properties"]
self._assert_system_prompt_captured(properties["$ai_input"])
def test_openai_separate_system_parameter(self):
"""Test OpenAI with system prompt as separate parameter."""
try:
from posthog.ai.openai import OpenAI
from openai.types.chat import ChatCompletion, ChatCompletionMessage
from openai.types.chat.chat_completion import Choice
from openai.types.completion_usage import CompletionUsage
from hanzo_insights.ai.openai import OpenAI
except ImportError:
self.skipTest("OpenAI package not available")
@@ -131,7 +137,7 @@ class TestSystemPromptCapture(unittest.TestCase):
"openai.resources.chat.completions.Completions.create",
return_value=mock_response,
):
client = OpenAI(posthog_client=self.client, api_key="test")
client = OpenAI(insights_client=self.client, api_key="test")
messages = [{"role": "user", "content": self.test_user_message}]
@@ -139,21 +145,24 @@ class TestSystemPromptCapture(unittest.TestCase):
model="gpt-4",
messages=messages,
system=self.test_system_prompt,
posthog_distinct_id="test-user",
insights_distinct_id="test-user",
)
self.assertEqual(len(self.client.capture.call_args_list), 1)
properties = self.client.capture.call_args_list[0][1]["properties"]
self.assertEqual(len(self.client._enqueue.call_args_list), 1)
properties = self.client._enqueue.call_args_list[0][0][0]["properties"]
self._assert_system_prompt_captured(properties["$ai_input"])
def test_openai_streaming_system_parameter(self):
"""Test OpenAI streaming with system parameter."""
try:
from posthog.ai.openai import OpenAI
from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
from openai.types.chat.chat_completion_chunk import (
ChatCompletionChunk,
ChoiceDelta,
)
from openai.types.chat.chat_completion_chunk import Choice as ChoiceChunk
from openai.types.chat.chat_completion_chunk import ChoiceDelta
from openai.types.completion_usage import CompletionUsage
from hanzo_insights.ai.openai import OpenAI
except ImportError:
self.skipTest("OpenAI package not available")
@@ -192,7 +201,7 @@ class TestSystemPromptCapture(unittest.TestCase):
"openai.resources.chat.completions.Completions.create",
return_value=[chunk1, chunk2],
):
client = OpenAI(posthog_client=self.client, api_key="test")
client = OpenAI(insights_client=self.client, api_key="test")
messages = [{"role": "user", "content": self.test_user_message}]
@@ -201,20 +210,20 @@ class TestSystemPromptCapture(unittest.TestCase):
messages=messages,
system=self.test_system_prompt,
stream=True,
posthog_distinct_id="test-user",
insights_distinct_id="test-user",
)
list(response_generator) # Consume generator
self.assertEqual(len(self.client.capture.call_args_list), 1)
properties = self.client.capture.call_args_list[0][1]["properties"]
self.assertEqual(len(self.client._enqueue.call_args_list), 1)
properties = self.client._enqueue.call_args_list[0][0][0]["properties"]
self._assert_system_prompt_captured(properties["$ai_input"])
# Anthropic Tests
def test_anthropic_messages_array_system_prompt(self):
"""Test Anthropic with system prompt in messages array."""
try:
from posthog.ai.anthropic import Anthropic
from hanzo_insights.ai.anthropic import Anthropic
except ImportError:
self.skipTest("Anthropic package not available")
@@ -226,7 +235,7 @@ class TestSystemPromptCapture(unittest.TestCase):
mock_response.usage.cache_creation_input_tokens = None
mock_create.return_value = mock_response
client = Anthropic(posthog_client=self.client, api_key="test")
client = Anthropic(insights_client=self.client, api_key="test")
messages = [
{"role": "system", "content": self.test_system_prompt},
@@ -236,17 +245,17 @@ class TestSystemPromptCapture(unittest.TestCase):
client.messages.create(
model="claude-3-5-sonnet-20241022",
messages=messages,
posthog_distinct_id="test-user",
insights_distinct_id="test-user",
)
self.assertEqual(len(self.client.capture.call_args_list), 1)
properties = self.client.capture.call_args_list[0][1]["properties"]
self.assertEqual(len(self.client._enqueue.call_args_list), 1)
properties = self.client._enqueue.call_args_list[0][0][0]["properties"]
self._assert_system_prompt_captured(properties["$ai_input"])
def test_anthropic_separate_system_parameter(self):
"""Test Anthropic with system prompt as separate parameter."""
try:
from posthog.ai.anthropic import Anthropic
from hanzo_insights.ai.anthropic import Anthropic
except ImportError:
self.skipTest("Anthropic package not available")
@@ -258,7 +267,7 @@ class TestSystemPromptCapture(unittest.TestCase):
mock_response.usage.cache_creation_input_tokens = None
mock_create.return_value = mock_response
client = Anthropic(posthog_client=self.client, api_key="test")
client = Anthropic(insights_client=self.client, api_key="test")
messages = [{"role": "user", "content": self.test_user_message}]
@@ -266,18 +275,18 @@ class TestSystemPromptCapture(unittest.TestCase):
model="claude-3-5-sonnet-20241022",
messages=messages,
system=self.test_system_prompt,
posthog_distinct_id="test-user",
insights_distinct_id="test-user",
)
self.assertEqual(len(self.client.capture.call_args_list), 1)
properties = self.client.capture.call_args_list[0][1]["properties"]
self.assertEqual(len(self.client._enqueue.call_args_list), 1)
properties = self.client._enqueue.call_args_list[0][0][0]["properties"]
self._assert_system_prompt_captured(properties["$ai_input"])
# Gemini Tests
def test_gemini_contents_array_system_prompt(self):
"""Test Gemini with system prompt in contents array."""
try:
from posthog.ai.gemini import Client
from hanzo_insights.ai.gemini import Client
except ImportError:
self.skipTest("Gemini package not available")
@@ -297,7 +306,7 @@ class TestSystemPromptCapture(unittest.TestCase):
mock_client_instance.models = mock_models_instance
mock_genai_class.return_value = mock_client_instance
client = Client(posthog_client=self.client, api_key="test")
client = Client(insights_client=self.client, api_key="test")
contents = [
{"role": "system", "content": self.test_system_prompt},
@@ -307,17 +316,17 @@ class TestSystemPromptCapture(unittest.TestCase):
client.models.generate_content(
model="gemini-2.0-flash",
contents=contents,
posthog_distinct_id="test-user",
insights_distinct_id="test-user",
)
self.assertEqual(len(self.client.capture.call_args_list), 1)
properties = self.client.capture.call_args_list[0][1]["properties"]
self.assertEqual(len(self.client._enqueue.call_args_list), 1)
properties = self.client._enqueue.call_args_list[0][0][0]["properties"]
self._assert_system_prompt_captured(properties["$ai_input"])
def test_gemini_system_instruction_parameter(self):
"""Test Gemini with system_instruction in config parameter."""
try:
from posthog.ai.gemini import Client
from hanzo_insights.ai.gemini import Client
except ImportError:
self.skipTest("Gemini package not available")
@@ -337,7 +346,7 @@ class TestSystemPromptCapture(unittest.TestCase):
mock_client_instance.models = mock_models_instance
mock_genai_class.return_value = mock_client_instance
client = Client(posthog_client=self.client, api_key="test")
client = Client(insights_client=self.client, api_key="test")
contents = [{"role": "user", "content": self.test_user_message}]
config = {"system_instruction": self.test_system_prompt}
@@ -346,9 +355,9 @@ class TestSystemPromptCapture(unittest.TestCase):
model="gemini-2.0-flash",
contents=contents,
config=config,
posthog_distinct_id="test-user",
insights_distinct_id="test-user",
)
self.assertEqual(len(self.client.capture.call_args_list), 1)
properties = self.client.capture.call_args_list[0][1]["properties"]
self.assertEqual(len(self.client._enqueue.call_args_list), 1)
properties = self.client._enqueue.call_args_list[0][0][0]["properties"]
self._assert_system_prompt_captured(properties["$ai_input"])
@@ -0,0 +1,62 @@
from parameterized import parameterized
from hanzo_insights.ai.utils import _get_tokens_source
@parameterized.expand(
[
("no_insights_properties", {"$ai_input_tokens": 100}, None, "sdk"),
("empty_insights_properties", {"$ai_input_tokens": 100}, {}, "sdk"),
(
"unrelated_insights_properties",
{"$ai_input_tokens": 100},
{"foo": "bar"},
"sdk",
),
(
"override_input_tokens",
{"$ai_input_tokens": 100},
{"$ai_input_tokens": 999},
"passthrough",
),
(
"override_output_tokens",
{"$ai_output_tokens": 50},
{"$ai_output_tokens": 999},
"passthrough",
),
(
"override_total_tokens",
{"$ai_input_tokens": 100},
{"$ai_total_tokens": 999},
"passthrough",
),
(
"override_cache_read",
{"$ai_input_tokens": 100},
{"$ai_cache_read_input_tokens": 500},
"passthrough",
),
(
"override_cache_creation",
{"$ai_input_tokens": 100},
{"$ai_cache_creation_input_tokens": 200},
"passthrough",
),
(
"override_reasoning_tokens",
{"$ai_input_tokens": 100},
{"$ai_reasoning_tokens": 300},
"passthrough",
),
(
"mixed_override_and_custom",
{"$ai_input_tokens": 100},
{"$ai_input_tokens": 999, "custom_key": "value"},
"passthrough",
),
]
)
def test_get_tokens_source(name, sdk_tags, insights_properties, expected):
result = _get_tokens_source(sdk_tags, insights_properties)
assert result == expected
@@ -1,4 +1,4 @@
from posthog.contexts import (
from hanzo_insights.contexts import (
new_context,
get_context_session_id,
get_context_distinct_id,
@@ -20,7 +20,7 @@ if not settings.configured:
)
django.setup()
from posthog.integrations.django import PosthogContextMiddleware
from hanzo_insights.integrations.django import InsightsContextMiddleware
class MockRequest:
@@ -45,7 +45,7 @@ class MockRequest:
return f"{scheme}://{self._host}{self.path}"
class TestPosthogContextMiddleware(unittest.TestCase):
class TestInsightsContextMiddleware(unittest.TestCase):
def create_middleware(
self,
extra_tags=None,
@@ -60,24 +60,24 @@ class TestPosthogContextMiddleware(unittest.TestCase):
with patch("django.conf.settings") as mock_settings:
# Configure mock settings
mock_settings.POSTHOG_MW_EXTRA_TAGS = extra_tags
mock_settings.POSTHOG_MW_REQUEST_FILTER = request_filter
mock_settings.POSTHOG_MW_TAG_MAP = tag_map
mock_settings.POSTHOG_MW_CAPTURE_EXCEPTIONS = capture_exceptions
mock_settings.POSTHOG_MW_CLIENT = None
mock_settings.INSIGHTS_MW_EXTRA_TAGS = extra_tags
mock_settings.INSIGHTS_MW_REQUEST_FILTER = request_filter
mock_settings.INSIGHTS_MW_TAG_MAP = tag_map
mock_settings.INSIGHTS_MW_CAPTURE_EXCEPTIONS = capture_exceptions
mock_settings.INSIGHTS_MW_CLIENT = None
# Make hasattr work correctly
def mock_hasattr(obj, name):
return name in [
"POSTHOG_MW_EXTRA_TAGS",
"POSTHOG_MW_REQUEST_FILTER",
"POSTHOG_MW_TAG_MAP",
"POSTHOG_MW_CAPTURE_EXCEPTIONS",
"POSTHOG_MW_CLIENT",
"INSIGHTS_MW_EXTRA_TAGS",
"INSIGHTS_MW_REQUEST_FILTER",
"INSIGHTS_MW_TAG_MAP",
"INSIGHTS_MW_CAPTURE_EXCEPTIONS",
"INSIGHTS_MW_CLIENT",
]
with patch("builtins.hasattr", side_effect=mock_hasattr):
middleware = PosthogContextMiddleware(get_response)
middleware = InsightsContextMiddleware(get_response)
return middleware
@@ -87,8 +87,8 @@ class TestPosthogContextMiddleware(unittest.TestCase):
middleware = self.create_middleware()
request = MockRequest(
headers={
"X-POSTHOG-SESSION-ID": "session-123",
"X-POSTHOG-DISTINCT-ID": "user-456",
"X-INSIGHTS-SESSION-ID": "session-123",
"X-INSIGHTS-DISTINCT-ID": "user-456",
},
method="POST",
path="/api/test",
@@ -104,7 +104,7 @@ class TestPosthogContextMiddleware(unittest.TestCase):
self.assertEqual(tags["$request_method"], "POST")
def test_extract_tags_missing_headers(self):
"""Test tag extraction when PostHog headers are missing"""
"""Test tag extraction when Insights headers are missing"""
with new_context():
middleware = self.create_middleware()
@@ -118,12 +118,12 @@ class TestPosthogContextMiddleware(unittest.TestCase):
self.assertEqual(tags["$request_method"], "GET")
def test_extract_tags_partial_headers(self):
"""Test tag extraction with only some PostHog headers present"""
"""Test tag extraction with only some Insights headers present"""
with new_context():
middleware = self.create_middleware()
request = MockRequest(
headers={"X-POSTHOG-SESSION-ID": "session-only"}, method="PUT"
headers={"X-INSIGHTS-SESSION-ID": "session-only"}, method="PUT"
)
tags = middleware.extract_tags(request)
@@ -141,7 +141,7 @@ class TestPosthogContextMiddleware(unittest.TestCase):
with new_context():
middleware = self.create_middleware(extra_tags=extra_tags_func)
request = MockRequest(
headers={"X-POSTHOG-SESSION-ID": "session-123"}, method="GET"
headers={"X-INSIGHTS-SESSION-ID": "session-123"}, method="GET"
)
tags = middleware.extract_tags(request)
@@ -167,7 +167,7 @@ class TestPosthogContextMiddleware(unittest.TestCase):
tag_map=tag_map_func, extra_tags=extra_tags_func
)
request = MockRequest(
headers={"X-POSTHOG-SESSION-ID": "session-123"}, method="GET"
headers={"X-INSIGHTS-SESSION-ID": "session-123"}, method="GET"
)
tags = middleware.extract_tags(request)
@@ -230,7 +230,7 @@ class TestPosthogContextMiddleware(unittest.TestCase):
middleware.client = mock_client
request = MockRequest(
headers={"X-POSTHOG-DISTINCT-ID": "test-user"},
headers={"X-INSIGHTS-DISTINCT-ID": "test-user"},
method="POST",
path="/api/endpoint",
)
@@ -282,7 +282,7 @@ class TestPosthogContextMiddleware(unittest.TestCase):
mock_client.capture_exception.assert_not_called()
class TestPosthogContextMiddlewareSync(unittest.TestCase):
class TestInsightsContextMiddlewareSync(unittest.TestCase):
"""Test synchronous middleware behavior"""
def test_sync_middleware_call(self):
@@ -291,13 +291,13 @@ class TestPosthogContextMiddlewareSync(unittest.TestCase):
get_response = Mock(return_value=mock_response)
# Create middleware with sync get_response
middleware = PosthogContextMiddleware(get_response)
middleware = InsightsContextMiddleware(get_response)
# Verify sync mode detected
self.assertFalse(middleware._is_coroutine)
request = MockRequest(
headers={"X-POSTHOG-SESSION-ID": "test-session"},
headers={"X-INSIGHTS-SESSION-ID": "test-session"},
method="GET",
path="/test",
)
@@ -315,8 +315,10 @@ class TestPosthogContextMiddlewareSync(unittest.TestCase):
get_response = Mock(return_value=mock_response)
# Create middleware with request filter that filters all requests
request_filter = lambda req: False
middleware = PosthogContextMiddleware.__new__(PosthogContextMiddleware)
def request_filter(req):
return False
middleware = InsightsContextMiddleware.__new__(InsightsContextMiddleware)
middleware.get_response = get_response
middleware._is_coroutine = False
middleware.request_filter = request_filter
@@ -349,7 +351,7 @@ class TestPosthogContextMiddlewareSync(unittest.TestCase):
mock_client = Mock()
get_response = Mock(return_value=Mock(status_code=500))
middleware = PosthogContextMiddleware(get_response)
middleware = InsightsContextMiddleware(get_response)
middleware.client = mock_client
def get_response_simulating_django(request):
@@ -378,7 +380,7 @@ class TestPosthogContextMiddlewareSync(unittest.TestCase):
)
class TestPosthogContextMiddlewareAsync(unittest.TestCase):
class TestInsightsContextMiddlewareAsync(unittest.TestCase):
"""Test asynchronous middleware behavior"""
def test_async_middleware_detection(self):
@@ -387,7 +389,7 @@ class TestPosthogContextMiddlewareAsync(unittest.TestCase):
async def async_get_response(request):
return Mock()
middleware = PosthogContextMiddleware(async_get_response)
middleware = InsightsContextMiddleware(async_get_response)
# Verify async mode detected
self.assertTrue(middleware._is_coroutine)
@@ -401,10 +403,10 @@ class TestPosthogContextMiddlewareAsync(unittest.TestCase):
async def async_get_response(request):
return mock_response
middleware = PosthogContextMiddleware(async_get_response)
middleware = InsightsContextMiddleware(async_get_response)
request = MockRequest(
headers={"X-POSTHOG-SESSION-ID": "async-session"},
headers={"X-INSIGHTS-SESSION-ID": "async-session"},
method="POST",
path="/async-test",
)
@@ -432,7 +434,7 @@ class TestPosthogContextMiddlewareAsync(unittest.TestCase):
return mock_response
# Properly initialize middleware
middleware = PosthogContextMiddleware(async_get_response)
middleware = InsightsContextMiddleware(async_get_response)
# Override request filter after initialization
middleware.request_filter = lambda req: False
@@ -457,10 +459,10 @@ class TestPosthogContextMiddlewareAsync(unittest.TestCase):
self.assertEqual(session_id, "async-session-123")
return mock_response
middleware = PosthogContextMiddleware(async_get_response)
middleware = InsightsContextMiddleware(async_get_response)
request = MockRequest(
headers={"X-POSTHOG-SESSION-ID": "async-session-123"},
headers={"X-INSIGHTS-SESSION-ID": "async-session-123"},
method="GET",
)
@@ -481,7 +483,7 @@ class TestPosthogContextMiddlewareAsync(unittest.TestCase):
raise ValueError("Async test exception")
# Properly initialize middleware
middleware = PosthogContextMiddleware(raise_exception)
middleware = InsightsContextMiddleware(raise_exception)
middleware.client = mock_client # Override with mock client
request = MockRequest()
@@ -523,11 +525,11 @@ class TestPosthogContextMiddlewareAsync(unittest.TestCase):
self.assertEqual(distinct_id, "123")
return mock_response
middleware = PosthogContextMiddleware(async_get_response)
middleware = InsightsContextMiddleware(async_get_response)
middleware.client = Mock()
request = MockRequest(
headers={"X-POSTHOG-SESSION-ID": "test-session"}, method="GET"
headers={"X-INSIGHTS-SESSION-ID": "test-session"}, method="GET"
)
# Mock auser() to return authenticated user
@@ -559,11 +561,11 @@ class TestPosthogContextMiddlewareAsync(unittest.TestCase):
self.assertIsNone(distinct_id)
return mock_response
middleware = PosthogContextMiddleware(async_get_response)
middleware = InsightsContextMiddleware(async_get_response)
middleware.client = Mock()
request = MockRequest(
headers={"X-POSTHOG-SESSION-ID": "test-session"}, method="GET"
headers={"X-INSIGHTS-SESSION-ID": "test-session"}, method="GET"
)
async def mock_auser():
@@ -589,12 +591,12 @@ class TestPosthogContextMiddlewareAsync(unittest.TestCase):
async def async_get_response(request):
return mock_response
middleware = PosthogContextMiddleware(async_get_response)
middleware = InsightsContextMiddleware(async_get_response)
middleware.client = Mock()
# Request without auser method (no auth middleware)
request = MockRequest(
headers={"X-POSTHOG-SESSION-ID": "test-session"}, method="GET"
headers={"X-INSIGHTS-SESSION-ID": "test-session"}, method="GET"
)
with new_context():
@@ -619,12 +621,12 @@ class TestPosthogContextMiddlewareAsync(unittest.TestCase):
async def async_get_response(request):
return mock_response
middleware = PosthogContextMiddleware(async_get_response)
middleware = InsightsContextMiddleware(async_get_response)
middleware.extra_tags = extra_tags_callback
middleware.client = Mock()
request = MockRequest(
headers={"X-POSTHOG-SESSION-ID": "test-session"}, method="GET"
headers={"X-INSIGHTS-SESSION-ID": "test-session"}, method="GET"
)
# Mock auser for no user
@@ -656,12 +658,12 @@ class TestPosthogContextMiddlewareAsync(unittest.TestCase):
async def async_get_response(request):
return mock_response
middleware = PosthogContextMiddleware(async_get_response)
middleware = InsightsContextMiddleware(async_get_response)
middleware.tag_map = tag_map_callback
middleware.client = Mock()
request = MockRequest(
headers={"X-POSTHOG-SESSION-ID": "test-session"}, method="GET"
headers={"X-INSIGHTS-SESSION-ID": "test-session"}, method="GET"
)
# Mock auser for no user
@@ -697,12 +699,12 @@ class TestPosthogContextMiddlewareAsync(unittest.TestCase):
self.assertEqual(session_id, "async-sess-123")
return mock_response
middleware = PosthogContextMiddleware(async_get_response)
middleware = InsightsContextMiddleware(async_get_response)
middleware.client = Mock()
request = MockRequest(
headers={
"X-POSTHOG-SESSION-ID": "async-sess-123",
"X-INSIGHTS-SESSION-ID": "async-sess-123",
"X-Forwarded-For": "192.168.1.1",
"User-Agent": "TestAgent/1.0",
},
@@ -723,13 +725,13 @@ class TestPosthogContextMiddlewareAsync(unittest.TestCase):
asyncio.run(run_test())
class TestPosthogContextMiddlewareHybrid(unittest.TestCase):
class TestInsightsContextMiddlewareHybrid(unittest.TestCase):
"""Test hybrid middleware behavior with mixed sync/async chains"""
def test_hybrid_flags_set(self):
"""Test that both capability flags are set"""
self.assertTrue(PosthogContextMiddleware.sync_capable)
self.assertTrue(PosthogContextMiddleware.async_capable)
self.assertTrue(InsightsContextMiddleware.sync_capable)
self.assertTrue(InsightsContextMiddleware.async_capable)
def test_sync_to_async_routing(self):
"""Test that __call__ routes to __acall__ when async"""
@@ -738,7 +740,7 @@ class TestPosthogContextMiddlewareHybrid(unittest.TestCase):
async def async_get_response(request):
return Mock()
middleware = PosthogContextMiddleware(async_get_response)
middleware = InsightsContextMiddleware(async_get_response)
# Verify routing happens
request = MockRequest()
@@ -757,7 +759,7 @@ class TestPosthogContextMiddlewareHybrid(unittest.TestCase):
def sync_get_response(request):
return mock_response
middleware = PosthogContextMiddleware(sync_get_response)
middleware = InsightsContextMiddleware(sync_get_response)
request = MockRequest()
result = middleware(request)
@@ -2,16 +2,16 @@ import unittest
import mock
from posthog.client import Client
from posthog.test.test_utils import FAKE_TEST_API_KEY
from hanzo_insights.client import Client
from hanzo_insights.test.test_utils import FAKE_TEST_API_KEY
class TestClient(unittest.TestCase):
@classmethod
def setUpClass(cls):
# This ensures no real HTTP POST requests are made
cls.client_post_patcher = mock.patch("posthog.client.batch_post")
cls.consumer_post_patcher = mock.patch("posthog.consumer.batch_post")
cls.client_post_patcher = mock.patch("hanzo_insights.client.batch_post")
cls.consumer_post_patcher = mock.patch("hanzo_insights.consumer.batch_post")
cls.client_post_patcher.start()
cls.consumer_post_patcher.start()
@@ -40,7 +40,7 @@ class TestClient(unittest.TestCase):
event["properties"]["processed_by_before_send"] = True
return event
with mock.patch("posthog.client.batch_post") as mock_post:
with mock.patch("hanzo_insights.client.batch_post") as mock_post:
client = Client(
FAKE_TEST_API_KEY,
on_error=self.set_fail,
@@ -73,7 +73,7 @@ class TestClient(unittest.TestCase):
return None
return event
with mock.patch("posthog.client.batch_post") as mock_post:
with mock.patch("hanzo_insights.client.batch_post") as mock_post:
client = Client(
FAKE_TEST_API_KEY,
on_error=self.set_fail,
@@ -101,7 +101,7 @@ class TestClient(unittest.TestCase):
def buggy_before_send(event):
raise ValueError("Oops!")
with mock.patch("posthog.client.batch_post") as mock_post:
with mock.patch("hanzo_insights.client.batch_post") as mock_post:
client = Client(
FAKE_TEST_API_KEY,
on_error=self.set_fail,
@@ -128,7 +128,7 @@ class TestClient(unittest.TestCase):
event["properties"]["marked"] = True
return event
with mock.patch("posthog.client.batch_post") as mock_post:
with mock.patch("hanzo_insights.client.batch_post") as mock_post:
client = Client(
FAKE_TEST_API_KEY,
on_error=self.set_fail,
@@ -153,7 +153,7 @@ class TestClient(unittest.TestCase):
def test_before_send_callback_disabled_when_none(self):
"""Test that client works normally when before_send is None."""
with mock.patch("posthog.client.batch_post") as mock_post:
with mock.patch("hanzo_insights.client.batch_post") as mock_post:
client = Client(
FAKE_TEST_API_KEY,
on_error=self.set_fail,
@@ -189,7 +189,7 @@ class TestClient(unittest.TestCase):
return event
with mock.patch("posthog.client.batch_post") as mock_post:
with mock.patch("hanzo_insights.client.batch_post") as mock_post:
client = Client(
FAKE_TEST_API_KEY,
on_error=self.set_fail,
File diff suppressed because it is too large Load Diff
+256
View File
@@ -0,0 +1,256 @@
import json
import time
import unittest
from typing import Any
import mock
from parameterized import parameterized
try:
from queue import Queue
except ImportError:
from Queue import Queue
from hanzo_insights.consumer import MAX_MSG_SIZE, Consumer
from hanzo_insights.request import APIError
from hanzo_insights.test.test_utils import TEST_API_KEY
def _track_event(event_name: str = "python event") -> dict[str, str]:
return {"type": "track", "event": event_name, "distinct_id": "distinct_id"}
class TestConsumer(unittest.TestCase):
def test_next(self) -> None:
q = Queue()
consumer = Consumer(q, "")
q.put(1)
next = consumer.next()
self.assertEqual(next, [1])
def test_next_limit(self) -> None:
q = Queue()
flush_at = 50
consumer = Consumer(q, "", flush_at)
for i in range(10000):
q.put(i)
next = consumer.next()
self.assertEqual(next, list(range(flush_at)))
def test_dropping_oversize_msg(self) -> None:
q = Queue()
consumer = Consumer(q, "")
oversize_msg = {"m": "x" * MAX_MSG_SIZE}
q.put(oversize_msg)
next = consumer.next()
self.assertEqual(next, [])
self.assertTrue(q.empty())
def test_upload(self) -> None:
q = Queue()
consumer = Consumer(q, TEST_API_KEY)
q.put(_track_event())
success = consumer.upload()
self.assertTrue(success)
def test_flush_interval(self) -> None:
# Put _n_ items in the queue, pausing a little bit more than
# _flush_interval_ after each one.
# The consumer should upload _n_ times.
q = Queue()
flush_interval = 0.3
consumer = Consumer(q, TEST_API_KEY, flush_at=10, flush_interval=flush_interval)
with mock.patch("hanzo_insights.consumer.batch_post") as mock_post:
consumer.start()
for i in range(3):
q.put(_track_event("python event %d" % i))
time.sleep(flush_interval * 1.1)
self.assertEqual(mock_post.call_count, 3)
def test_multiple_uploads_per_interval(self) -> None:
# Put _flush_at*2_ items in the queue at once, then pause for
# _flush_interval_. The consumer should upload 2 times.
q = Queue()
flush_interval = 0.5
flush_at = 10
consumer = Consumer(
q, TEST_API_KEY, flush_at=flush_at, flush_interval=flush_interval
)
with mock.patch("hanzo_insights.consumer.batch_post") as mock_post:
consumer.start()
for i in range(flush_at * 2):
q.put(_track_event("python event %d" % i))
time.sleep(flush_interval * 1.1)
self.assertEqual(mock_post.call_count, 2)
def test_request(self) -> None:
consumer = Consumer(None, TEST_API_KEY)
consumer.request([_track_event()])
def _run_retry_test(
self, exception: Exception, exception_count: int, retries: int = 10
) -> None:
call_count = [0]
def mock_post(*args: Any, **kwargs: Any) -> None:
call_count[0] += 1
if call_count[0] <= exception_count:
raise exception
consumer = Consumer(None, TEST_API_KEY, retries=retries)
with mock.patch(
"hanzo_insights.consumer.batch_post", mock.Mock(side_effect=mock_post)
):
if exception_count <= retries:
consumer.request([_track_event()])
else:
with self.assertRaises(type(exception)):
consumer.request([_track_event()])
@parameterized.expand(
[
("general_errors", Exception("generic exception"), 2),
("server_errors", APIError(500, "Internal Server Error"), 2),
("rate_limit_errors", APIError(429, "Too Many Requests"), 2),
]
)
def test_request_retries_on_retriable_errors(
self, _name: str, exception: Exception, exception_count: int
) -> None:
self._run_retry_test(exception, exception_count)
def test_request_does_not_retry_client_errors(self) -> None:
with self.assertRaises(APIError):
self._run_retry_test(APIError(400, "Client Errors"), 1)
def test_request_fails_when_exceptions_exceed_retries(self) -> None:
self._run_retry_test(APIError(500, "Internal Server Error"), 4, retries=3)
def test_pause(self) -> None:
consumer = Consumer(None, TEST_API_KEY)
consumer.pause()
self.assertFalse(consumer.running)
def test_max_batch_size(self) -> None:
q = Queue()
consumer = Consumer(q, TEST_API_KEY, flush_at=100000, flush_interval=3)
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())
# Let's capture 8MB of data to trigger two batches
n_msgs = int(8_000_000 / msg_size)
def mock_post_fn(_: str, data: str, **kwargs: Any) -> mock.Mock:
res = mock.Mock()
res.status_code = 200
request_size = len(data.encode())
# Batches close after the first message bringing it bigger than BATCH_SIZE_LIMIT, let's add 10% of margin
self.assertTrue(
request_size < (5 * 1024 * 1024) * 1.1,
"batch size (%d) higher than limit" % request_size,
)
return res
with mock.patch(
"hanzo_insights.request._session.post", side_effect=mock_post_fn
) as mock_post:
consumer.start()
for _ in range(0, n_msgs + 2):
q.put(track)
q.join()
self.assertEqual(mock_post.call_count, 2)
def test_request_sleeps_with_retry_after(self) -> None:
error = APIError(429, "Too Many Requests", retry_after=5.0)
call_count = [0]
def mock_post(*args: Any, **kwargs: Any) -> None:
call_count[0] += 1
if call_count[0] <= 1:
raise error
consumer = Consumer(None, TEST_API_KEY, retries=3)
with (
mock.patch("hanzo_insights.consumer.batch_post", side_effect=mock_post),
mock.patch("hanzo_insights.consumer.time.sleep") as mock_sleep,
):
consumer.request([_track_event()])
mock_sleep.assert_called_once_with(5.0)
def test_request_uses_exponential_backoff_without_retry_after(self) -> None:
error = APIError(503, "Service Unavailable")
call_count = [0]
def mock_post(*args: Any, **kwargs: Any) -> None:
call_count[0] += 1
if call_count[0] <= 3:
raise error
consumer = Consumer(None, TEST_API_KEY, retries=3)
with (
mock.patch("hanzo_insights.consumer.batch_post", side_effect=mock_post),
mock.patch("hanzo_insights.consumer.time.sleep") as mock_sleep,
):
consumer.request([_track_event()])
self.assertEqual(
mock_sleep.call_args_list,
[
mock.call(1), # 2^0
mock.call(2), # 2^1
mock.call(4), # 2^2
],
)
def test_request_retries_on_408(self) -> None:
call_count = [0]
def mock_post(*args: Any, **kwargs: Any) -> None:
call_count[0] += 1
if call_count[0] <= 1:
raise APIError(408, "Request Timeout")
consumer = Consumer(None, TEST_API_KEY, retries=3)
with (
mock.patch("hanzo_insights.consumer.batch_post", side_effect=mock_post),
mock.patch("hanzo_insights.consumer.time.sleep"),
):
consumer.request([_track_event()])
self.assertEqual(call_count[0], 2)
@parameterized.expand(
[
("on_error_succeeds", False),
("on_error_raises", True),
]
)
def test_upload_exception_calls_on_error_and_does_not_raise(
self, _name: str, on_error_raises: bool
) -> None:
on_error_called: list[tuple[Exception, list[dict[str, str]]]] = []
def on_error(e: Exception, batch: list[dict[str, str]]) -> None:
on_error_called.append((e, batch))
if on_error_raises:
raise Exception("on_error failed")
q = Queue()
consumer = Consumer(q, TEST_API_KEY, on_error=on_error)
track = _track_event()
q.put(track)
with mock.patch.object(
consumer, "request", side_effect=Exception("request failed")
):
result = consumer.upload()
self.assertFalse(result)
self.assertEqual(len(on_error_called), 1)
self.assertEqual(str(on_error_called[0][0]), "request failed")
self.assertEqual(on_error_called[0][1], [track])
@@ -1,7 +1,7 @@
import unittest
from unittest.mock import patch
from posthog.contexts import (
from hanzo_insights.contexts import (
get_tags,
new_context,
scoped,
@@ -66,7 +66,7 @@ class TestContexts(unittest.TestCase):
# Back to level 1
assert get_tags() == {"level1": "value1"}
@patch("posthog.capture_exception")
@patch("hanzo_insights.capture_exception")
def test_scoped_decorator_success(self, mock_capture):
@scoped()
def successful_function(x, y):
@@ -85,7 +85,7 @@ class TestContexts(unittest.TestCase):
# Context should be cleared after function execution
assert get_tags() == {}
@patch("posthog.capture_exception")
@patch("hanzo_insights.capture_exception")
def test_scoped_decorator_exception(self, mock_capture):
test_exception = ValueError("Test exception")
@@ -111,7 +111,7 @@ class TestContexts(unittest.TestCase):
# Context should be cleared after function execution
assert get_tags() == {}
@patch("posthog.capture_exception")
@patch("hanzo_insights.capture_exception")
def test_new_context_exception_handling(self, mock_capture):
test_exception = RuntimeError("Context exception")
@@ -191,6 +191,32 @@ class TestContexts(unittest.TestCase):
assert get_context_distinct_id() == "user123"
assert get_context_session_id() == "session456"
def test_child_tags_override_parent_tags_in_non_fresh_context(self):
with new_context(fresh=True):
tag("shared_key", "parent_value")
tag("parent_only", "parent")
with new_context(fresh=False):
# Child should inherit parent tags
assert get_tags()["parent_only"] == "parent"
# Child sets same key - should override parent
tag("shared_key", "child_value")
tag("child_only", "child")
tags = get_tags()
# Child value should win for shared key
assert tags["shared_key"] == "child_value"
# Both parent and child tags should be present
assert tags["parent_only"] == "parent"
assert tags["child_only"] == "child"
# Parent context should be unchanged
parent_tags = get_tags()
assert parent_tags["shared_key"] == "parent_value"
assert parent_tags["parent_only"] == "parent"
assert "child_only" not in parent_tags
def test_scoped_decorator_with_context_ids(self):
@scoped()
def function_with_context():
@@ -0,0 +1,689 @@
import subprocess
import sys
from textwrap import dedent
import pytest
def test_excepthook(tmpdir):
app = tmpdir.join("app.py")
app.write(
dedent(
"""
from hanzo_insights import Insights
client = Insights('phc_x', host='https://eu.i.insights.hanzo.ai', 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:hanzo_insights: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_code_variables_capture(tmpdir):
app = tmpdir.join("app.py")
app.write(
dedent(
"""
import os
from hanzo_insights import Insights
class UnserializableObject:
pass
client = Insights(
'phc_x',
host='https://eu.i.insights.hanzo.ai',
debug=True,
enable_exception_autocapture=True,
capture_exception_code_variables=True,
project_root=os.path.dirname(os.path.abspath(__file__))
)
def trigger_error():
my_string = "hello world"
my_number = 42
my_bool = True
my_dict = {"name": "test", "value": 123}
my_sensitive_dict = {
"safe_key": "safe_value",
"password": "secret123", # key matches pattern -> should be masked
"other_key": "contains_password_here", # value matches pattern -> should be masked
}
my_nested_dict = {
"level1": {
"level2": {
"api_key": "nested_secret", # deeply nested key matches
"data": "contains_token_here", # deeply nested value matches
"safe": "visible",
}
}
}
my_list = ["safe_item", "has_password_inside", "another_safe"]
my_tuple = ("tuple_safe", "secret_in_value", "tuple_also_safe")
my_list_of_dicts = [
{"id": 1, "password": "list_dict_secret"},
{"id": 2, "value": "safe_value"},
]
my_obj = UnserializableObject()
my_password = "secret123" # Should be masked by default (name matches)
my_innocent_var = "contains_password_here" # Should be masked by default (value matches)
__should_be_ignored = "hidden" # Should be ignored by default
1/0 # Trigger exception
def intermediate_function():
request_id = "abc-123"
user_count = 100
is_active = True
trigger_error()
def process_data():
batch_size = 50
retry_count = 3
intermediate_function()
process_data()
"""
)
)
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"code_variables" in output
# Variables from trigger_error frame
assert b"'my_string': 'hello world'" in output
assert b"'my_number': 42" in output
assert b"'my_bool': 'True'" in output
assert b'"my_dict": "{\\"name\\": \\"test\\", \\"value\\": 123}"' in output
assert (
b'{\\"safe_key\\": \\"safe_value\\", \\"password\\": \\"$$_insights_redacted_based_on_masking_rules_$$\\", \\"other_key\\": \\"$$_insights_redacted_based_on_masking_rules_$$\\"}'
in output
)
assert (
b'{\\"level1\\": {\\"level2\\": {\\"api_key\\": \\"$$_insights_redacted_based_on_masking_rules_$$\\", \\"data\\": \\"$$_insights_redacted_based_on_masking_rules_$$\\", \\"safe\\": \\"visible\\"}}}'
in output
)
assert (
b'[\\"safe_item\\", \\"$$_insights_redacted_based_on_masking_rules_$$\\", \\"another_safe\\"]'
in output
)
assert (
b'[\\"tuple_safe\\", \\"$$_insights_redacted_based_on_masking_rules_$$\\", \\"tuple_also_safe\\"]'
in output
)
assert (
b'[{\\"id\\": 1, \\"password\\": \\"$$_insights_redacted_based_on_masking_rules_$$\\"}, {\\"id\\": 2, \\"value\\": \\"safe_value\\"}]'
in output
)
assert b"<__main__.UnserializableObject object at" in output
assert b"'my_password': '$$_insights_redacted_based_on_masking_rules_$$'" in output
assert (
b"'my_innocent_var': '$$_insights_redacted_based_on_masking_rules_$$'" in output
)
assert b"'__should_be_ignored':" not in output
# Variables from intermediate_function frame
assert b"'request_id': 'abc-123'" in output
assert b"'user_count': 100" in output
assert b"'is_active': 'True'" in output
# Variables from process_data frame
assert b"'batch_size': 50" in output
assert b"'retry_count': 3" in output
def test_code_variables_context_override(tmpdir):
app = tmpdir.join("app.py")
app.write(
dedent(
"""
import os
import hanzo_insights
from hanzo_insights import Insights
insights_client = Insights(
'phc_x',
host='https://eu.i.insights.hanzo.ai',
debug=True,
enable_exception_autocapture=True,
capture_exception_code_variables=False,
project_root=os.path.dirname(os.path.abspath(__file__))
)
def process_data():
bank = "should_be_masked"
__dunder_var = "should_be_visible"
1/0
with hanzo_insights.new_context(client=insights_client):
hanzo_insights.set_capture_exception_code_variables_context(True)
hanzo_insights.set_code_variables_mask_patterns_context([r"(?i).*bank.*"])
hanzo_insights.set_code_variables_ignore_patterns_context([])
process_data()
"""
)
)
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"code_variables" in output
assert b"'bank': '$$_insights_redacted_based_on_masking_rules_$$'" in output
assert b"'__dunder_var': 'should_be_visible'" in output
def test_code_variables_size_limiter(tmpdir):
app = tmpdir.join("app.py")
app.write(
dedent(
"""
import os
from hanzo_insights import Insights
client = Insights(
'phc_x',
host='https://eu.i.insights.hanzo.ai',
debug=True,
enable_exception_autocapture=True,
capture_exception_code_variables=True,
project_root=os.path.dirname(os.path.abspath(__file__))
)
def trigger_error():
var_a = "a" * 2000
var_b = "b" * 2000
var_c = "c" * 2000
var_d = "d" * 2000
var_e = "e" * 2000
var_f = "f" * 2000
var_g = "g" * 2000
1/0
def intermediate_function():
var_h = "h" * 2000
var_i = "i" * 2000
var_j = "j" * 2000
var_k = "k" * 2000
var_l = "l" * 2000
var_m = "m" * 2000
var_n = "n" * 2000
trigger_error()
def process_data():
var_o = "o" * 2000
var_p = "p" * 2000
var_q = "q" * 2000
var_r = "r" * 2000
var_s = "s" * 2000
var_t = "t" * 2000
var_u = "u" * 2000
intermediate_function()
process_data()
"""
)
)
with pytest.raises(subprocess.CalledProcessError) as excinfo:
subprocess.check_output([sys.executable, str(app)], stderr=subprocess.STDOUT)
output = excinfo.value.output.decode("utf-8")
assert "ZeroDivisionError" in output
assert "code_variables" in output
captured_vars = []
for var_name in [
"var_a",
"var_b",
"var_c",
"var_d",
"var_e",
"var_f",
"var_g",
"var_h",
"var_i",
"var_j",
"var_k",
"var_l",
"var_m",
"var_n",
"var_o",
"var_p",
"var_q",
"var_r",
"var_s",
"var_t",
"var_u",
]:
if f"'{var_name}'" in output:
captured_vars.append(var_name)
assert len(captured_vars) > 0
assert len(captured_vars) < 21
def test_code_variables_disabled_capture(tmpdir):
app = tmpdir.join("app.py")
app.write(
dedent(
"""
import os
from hanzo_insights import Insights
client = Insights(
'phc_x',
host='https://eu.i.insights.hanzo.ai',
debug=True,
enable_exception_autocapture=True,
capture_exception_code_variables=False,
project_root=os.path.dirname(os.path.abspath(__file__))
)
def trigger_error():
my_string = "hello world"
my_number = 42
my_bool = True
1/0
trigger_error()
"""
)
)
with pytest.raises(subprocess.CalledProcessError) as excinfo:
subprocess.check_output([sys.executable, str(app)], stderr=subprocess.STDOUT)
output = excinfo.value.output.decode("utf-8")
assert "ZeroDivisionError" in output
assert "'code_variables':" not in output
assert '"code_variables":' not in output
assert "'my_string'" not in output
assert "'my_number'" not in output
def test_code_variables_enabled_then_disabled_in_context(tmpdir):
app = tmpdir.join("app.py")
app.write(
dedent(
"""
import os
import hanzo_insights
from hanzo_insights import Insights
insights_client = Insights(
'phc_x',
host='https://eu.i.insights.hanzo.ai',
debug=True,
enable_exception_autocapture=True,
capture_exception_code_variables=True,
project_root=os.path.dirname(os.path.abspath(__file__))
)
def process_data():
my_var = "should not be captured"
important_value = 123
1/0
with hanzo_insights.new_context(client=insights_client):
hanzo_insights.set_capture_exception_code_variables_context(False)
process_data()
"""
)
)
with pytest.raises(subprocess.CalledProcessError) as excinfo:
subprocess.check_output([sys.executable, str(app)], stderr=subprocess.STDOUT)
output = excinfo.value.output.decode("utf-8")
assert "ZeroDivisionError" in output
assert "'code_variables':" not in output
assert '"code_variables":' not in output
assert "'my_var'" not in output
assert "'important_value'" not in output
def test_code_variables_repr_fallback(tmpdir):
app = tmpdir.join("app.py")
app.write(
dedent(
"""
import os
import re
from datetime import datetime, timedelta
from decimal import Decimal
from fractions import Fraction
from hanzo_insights import Insights
class CustomReprClass:
def __repr__(self):
return '<CustomReprClass: custom representation>'
client = Insights(
'phc_x',
host='https://eu.i.insights.hanzo.ai',
debug=True,
enable_exception_autocapture=True,
capture_exception_code_variables=True,
project_root=os.path.dirname(os.path.abspath(__file__))
)
def trigger_error():
my_regex = re.compile(r'\\d+')
my_datetime = datetime(2024, 1, 15, 10, 30, 45)
my_timedelta = timedelta(days=5, hours=3)
my_decimal = Decimal('123.456')
my_fraction = Fraction(3, 4)
my_set = {1, 2, 3}
my_frozenset = frozenset([4, 5, 6])
my_bytes = b'hello bytes'
my_bytearray = bytearray(b'mutable bytes')
my_memoryview = memoryview(b'memory view')
my_complex = complex(3, 4)
my_range = range(10)
my_custom = CustomReprClass()
my_lambda = lambda x: x * 2
my_function = trigger_error
1/0
trigger_error()
"""
)
)
with pytest.raises(subprocess.CalledProcessError) as excinfo:
subprocess.check_output([sys.executable, str(app)], stderr=subprocess.STDOUT)
output = excinfo.value.output.decode("utf-8")
assert "ZeroDivisionError" in output
assert "code_variables" in output
assert "re.compile(" in output and "\\\\d+" in output
assert "datetime.datetime(2024, 1, 15, 10, 30, 45)" in output
assert "datetime.timedelta(days=5, seconds=10800)" in output
assert "Decimal('123.456')" in output
assert "Fraction(3, 4)" in output
assert "{1, 2, 3}" in output
assert "frozenset({4, 5, 6})" in output
assert "b'hello bytes'" in output
assert "bytearray(b'mutable bytes')" in output
assert "<memory at" in output
assert "(3+4j)" in output
assert "range(0, 10)" in output
assert "<CustomReprClass: custom representation>" in output
assert "<lambda>" in output
assert "<function trigger_error at" in output
def test_code_variables_too_long_string_value_replaced(tmpdir):
app = tmpdir.join("app.py")
app.write(
dedent(
"""
import os
from hanzo_insights import Insights
client = Insights(
'phc_x',
host='https://eu.i.insights.hanzo.ai',
debug=True,
enable_exception_autocapture=True,
capture_exception_code_variables=True,
project_root=os.path.dirname(os.path.abspath(__file__))
)
def trigger_error():
short_value = "I am short"
long_value = "x" * 20000
long_blob = "password_" + "a" * 20000
1/0
trigger_error()
"""
)
)
with pytest.raises(subprocess.CalledProcessError) as excinfo:
subprocess.check_output([sys.executable, str(app)], stderr=subprocess.STDOUT)
output = excinfo.value.output.decode("utf-8")
assert "ZeroDivisionError" in output
assert "code_variables" in output
assert "'short_value': 'I am short'" in output
assert "$$_insights_value_too_long_$$" in output
assert "'long_blob': '$$_insights_value_too_long_$$'" in output
def test_code_variables_too_long_string_in_nested_dict(tmpdir):
app = tmpdir.join("app.py")
app.write(
dedent(
"""
import os
from hanzo_insights import Insights
client = Insights(
'phc_x',
host='https://eu.i.insights.hanzo.ai',
debug=True,
enable_exception_autocapture=True,
capture_exception_code_variables=True,
project_root=os.path.dirname(os.path.abspath(__file__))
)
def trigger_error():
my_data = {
"short_key": "short_val",
"long_key": "y" * 20000,
"nested": {
"deep_long": "z" * 20000,
"deep_short": "ok",
},
}
1/0
trigger_error()
"""
)
)
with pytest.raises(subprocess.CalledProcessError) as excinfo:
subprocess.check_output([sys.executable, str(app)], stderr=subprocess.STDOUT)
output = excinfo.value.output.decode("utf-8")
assert "ZeroDivisionError" in output
assert "code_variables" in output
assert "short_val" in output
assert "ok" in output
assert "$$_insights_value_too_long_$$" in output
assert "y" * 1000 not in output
assert "z" * 1000 not in output
def test_mask_sensitive_data_too_long_dict_key():
from hanzo_insights.exception_utils import (
CODE_VARIABLES_TOO_LONG_VALUE,
_compile_patterns,
_mask_sensitive_data,
)
compiled_mask = _compile_patterns([r"(?i)password"])
result = _mask_sensitive_data(
{
"short": "visible",
"k" * 20000: "hidden_val",
"password": "secret",
},
compiled_mask,
)
assert result["short"] == "visible"
# This then gets shortened by the JSON truncation at 1024 chars anyways so no worries
assert result["k" * 20000] == CODE_VARIABLES_TOO_LONG_VALUE
assert result["password"] == "$$_insights_redacted_based_on_masking_rules_$$"
def test_mask_sensitive_data_circular_ref():
from hanzo_insights.exception_utils import _compile_patterns, _mask_sensitive_data
compiled_mask = _compile_patterns([r"(?i)password"])
# Circular dict
circular_dict = {"key": "value"}
circular_dict["self"] = circular_dict
result = _mask_sensitive_data(circular_dict, compiled_mask)
assert result["key"] == "value"
assert result["self"] == "<circular ref>"
# Circular list
circular_list = ["item"]
circular_list.append(circular_list)
result = _mask_sensitive_data(circular_list, compiled_mask)
assert result[0] == "item"
assert result[1] == "<circular ref>"
def test_compile_patterns_fast_path_and_regex_fallback():
from hanzo_insights.exception_utils import _compile_patterns, _pattern_matches
# Simple case-insensitive patterns should become substrings
simple_only = _compile_patterns([r"(?i)password", r"(?i)token", r"(?i)jwt"])
substrings, regexes = simple_only
assert substrings == ["password", "token", "jwt"]
assert regexes == []
assert _pattern_matches("my_password_var", simple_only) is True
assert _pattern_matches("MY_TOKEN", simple_only) is True
assert _pattern_matches("safe_variable", simple_only) is False
# Complex regex patterns should stay as compiled regexes
complex_only = _compile_patterns([r"^__.*", r"\d{3,}", r"^sk_live_"])
substrings, regexes = complex_only
assert substrings == []
assert len(regexes) == 3
assert _pattern_matches("__dunder", complex_only) is True
assert _pattern_matches("has_999_numbers", complex_only) is True
assert _pattern_matches("sk_live_abc123", complex_only) is True
assert _pattern_matches("normal_var", complex_only) is False
# Mixed: simple substrings + complex regexes together
mixed = _compile_patterns(
[
r"(?i)secret", # simple
r"(?i)api_key", # simple
r"^__.*", # regex
r"\btoken_\w+", # regex
]
)
substrings, regexes = mixed
assert substrings == ["secret", "api_key"]
assert len(regexes) == 2
# Substring matches
assert _pattern_matches("my_secret", mixed) is True
assert _pattern_matches("API_KEY_VALUE", mixed) is True
# Regex matches
assert _pattern_matches("__private", mixed) is True
assert _pattern_matches("token_abc", mixed) is True
# No match
assert _pattern_matches("safe_var", mixed) is False
def test_mask_sensitive_data_large_dict_replaced():
from hanzo_insights.exception_utils import (
CODE_VARIABLES_TOO_LONG_VALUE,
_compile_patterns,
_mask_sensitive_data,
)
compiled_mask = _compile_patterns([r"(?i)password"])
large_dict = {f"key_{i}": f"value_{i}" for i in range(300)}
result = _mask_sensitive_data(large_dict, compiled_mask)
assert result == CODE_VARIABLES_TOO_LONG_VALUE
def test_mask_sensitive_data_large_list_replaced():
from hanzo_insights.exception_utils import (
CODE_VARIABLES_TOO_LONG_VALUE,
_compile_patterns,
_mask_sensitive_data,
)
compiled_mask = _compile_patterns([r"(?i)password"])
large_list = [f"item_{i}" for i in range(300)]
result = _mask_sensitive_data(large_list, compiled_mask)
assert result == CODE_VARIABLES_TOO_LONG_VALUE
def test_mask_sensitive_data_large_tuple_replaced():
from hanzo_insights.exception_utils import (
CODE_VARIABLES_TOO_LONG_VALUE,
_compile_patterns,
_mask_sensitive_data,
)
compiled_mask = _compile_patterns([r"(?i)password"])
large_tuple = tuple(f"item_{i}" for i in range(300))
result = _mask_sensitive_data(large_tuple, compiled_mask)
assert result == CODE_VARIABLES_TOO_LONG_VALUE
@@ -1,6 +1,6 @@
import unittest
from posthog.types import FeatureFlag, FlagMetadata, FlagReason, LegacyFlagMetadata
from hanzo_insights.types import FeatureFlag, FlagMetadata, FlagReason, LegacyFlagMetadata
class TestFeatureFlag(unittest.TestCase):
@@ -0,0 +1,883 @@
import unittest
import mock
from hanzo_insights.client import Client
from hanzo_insights.test.test_utils import FAKE_TEST_API_KEY
from hanzo_insights.types import (
FeatureFlag,
FeatureFlagError,
FeatureFlagResult,
FlagMetadata,
FlagReason,
)
class TestFeatureFlagResult(unittest.TestCase):
def test_from_bool_value_and_payload(self):
result = FeatureFlagResult.from_value_and_payload(
"test-flag", True, "[1, 2, 3]"
)
self.assertEqual(result.key, "test-flag")
self.assertEqual(result.enabled, True)
self.assertEqual(result.variant, None)
self.assertEqual(result.payload, [1, 2, 3])
def test_from_false_value_and_payload(self):
result = FeatureFlagResult.from_value_and_payload(
"test-flag", False, '{"some": "value"}'
)
self.assertEqual(result.key, "test-flag")
self.assertEqual(result.enabled, False)
self.assertEqual(result.variant, None)
self.assertEqual(result.payload, {"some": "value"})
def test_from_variant_value_and_payload(self):
result = FeatureFlagResult.from_value_and_payload(
"test-flag", "control", "true"
)
self.assertEqual(result.key, "test-flag")
self.assertEqual(result.enabled, True)
self.assertEqual(result.variant, "control")
self.assertEqual(result.payload, True)
def test_from_none_value_and_payload(self):
result = FeatureFlagResult.from_value_and_payload(
"test-flag", None, '{"some": "value"}'
)
self.assertIsNone(result)
def test_from_boolean_flag_details(self):
flag_details = FeatureFlag(
key="test-flag",
enabled=True,
variant=None,
metadata=FlagMetadata(
id=1, version=1, description="test-flag", payload='"Some string"'
),
reason=FlagReason(
code="test-reason", description="test-reason", condition_index=0
),
)
result = FeatureFlagResult.from_flag_details(flag_details)
self.assertEqual(result.key, "test-flag")
self.assertEqual(result.enabled, True)
self.assertEqual(result.variant, None)
self.assertEqual(result.payload, "Some string")
def test_from_boolean_flag_details_with_override_variant_match_value(self):
flag_details = FeatureFlag(
key="test-flag",
enabled=True,
variant=None,
metadata=FlagMetadata(
id=1, version=1, description="test-flag", payload='"Some string"'
),
reason=FlagReason(
code="test-reason", description="test-reason", condition_index=0
),
)
result = FeatureFlagResult.from_flag_details(
flag_details, override_match_value="control"
)
self.assertEqual(result.key, "test-flag")
self.assertEqual(result.enabled, True)
self.assertEqual(result.variant, "control")
self.assertEqual(result.payload, "Some string")
def test_from_boolean_flag_details_with_override_boolean_match_value(self):
flag_details = FeatureFlag(
key="test-flag",
enabled=True,
variant="control",
metadata=FlagMetadata(
id=1, version=1, description="test-flag", payload='{"some": "value"}'
),
reason=FlagReason(
code="test-reason", description="test-reason", condition_index=0
),
)
result = FeatureFlagResult.from_flag_details(
flag_details, override_match_value=True
)
self.assertEqual(result.key, "test-flag")
self.assertEqual(result.enabled, True)
self.assertEqual(result.variant, None)
self.assertEqual(result.payload, {"some": "value"})
def test_from_boolean_flag_details_with_override_false_match_value(self):
flag_details = FeatureFlag(
key="test-flag",
enabled=True,
variant="control",
metadata=FlagMetadata(
id=1, version=1, description="test-flag", payload='{"some": "value"}'
),
reason=FlagReason(
code="test-reason", description="test-reason", condition_index=0
),
)
result = FeatureFlagResult.from_flag_details(
flag_details, override_match_value=False
)
self.assertEqual(result.key, "test-flag")
self.assertEqual(result.enabled, False)
self.assertEqual(result.variant, None)
self.assertEqual(result.payload, {"some": "value"})
def test_from_variant_flag_details(self):
flag_details = FeatureFlag(
key="test-flag",
enabled=True,
variant="control",
metadata=FlagMetadata(
id=1, version=1, description="test-flag", payload='{"some": "value"}'
),
reason=FlagReason(
code="test-reason", description="test-reason", condition_index=0
),
)
result = FeatureFlagResult.from_flag_details(flag_details)
self.assertEqual(result.key, "test-flag")
self.assertEqual(result.enabled, True)
self.assertEqual(result.variant, "control")
self.assertEqual(result.payload, {"some": "value"})
def test_from_none_flag_details(self):
result = FeatureFlagResult.from_flag_details(None)
self.assertIsNone(result)
def test_from_flag_details_with_none_payload(self):
flag_details = FeatureFlag(
key="test-flag",
enabled=True,
variant=None,
metadata=FlagMetadata(
id=1, version=1, description="test-flag", payload=None
),
reason=FlagReason(
code="test-reason", description="test-reason", condition_index=0
),
)
result = FeatureFlagResult.from_flag_details(flag_details)
self.assertEqual(result.key, "test-flag")
self.assertEqual(result.enabled, True)
self.assertEqual(result.variant, None)
self.assertIsNone(result.payload)
class TestGetFeatureFlagResult(unittest.TestCase):
@classmethod
def setUpClass(cls):
# This ensures no real HTTP POST requests are made
cls.capture_patch = mock.patch.object(Client, "capture")
cls.capture_patch.start()
@classmethod
def tearDownClass(cls):
cls.capture_patch.stop()
def set_fail(self, e, batch):
"""Mark the failure handler"""
self.failed = True
def setUp(self):
self.failed = False
self.client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail)
@mock.patch.object(Client, "capture")
def test_get_feature_flag_result_boolean_local_evaluation(self, patch_capture):
basic_flag = {
"id": 1,
"name": "Beta Feature",
"key": "person-flag",
"active": True,
"filters": {
"groups": [
{
"properties": [
{
"key": "region",
"operator": "exact",
"value": ["USA"],
"type": "person",
}
],
"rollout_percentage": 100,
}
],
"payloads": {"true": "300"},
},
}
self.client.feature_flags = [basic_flag]
flag_result = self.client.get_feature_flag_result(
"person-flag", "some-distinct-id", person_properties={"region": "USA"}
)
self.assertEqual(flag_result.enabled, True)
self.assertEqual(flag_result.variant, None)
self.assertEqual(flag_result.payload, 300)
patch_capture.assert_called_with(
"$feature_flag_called",
distinct_id="some-distinct-id",
properties={
"$feature_flag": "person-flag",
"$feature_flag_response": True,
"locally_evaluated": True,
"$feature/person-flag": True,
"$feature_flag_payload": 300,
},
groups={},
disable_geoip=None,
)
# Verify error property is NOT present on successful evaluation
captured_properties = patch_capture.call_args[1]["properties"]
self.assertNotIn("$feature_flag_error", captured_properties)
@mock.patch.object(Client, "capture")
def test_get_feature_flag_result_variant_local_evaluation(self, patch_capture):
basic_flag = {
"id": 1,
"name": "Beta Feature",
"key": "person-flag",
"active": True,
"filters": {
"groups": [
{
"properties": [
{
"key": "region",
"operator": "exact",
"value": ["USA"],
"type": "person",
}
],
"rollout_percentage": 100,
}
],
"multivariate": {
"variants": [
{"key": "variant-1", "rollout_percentage": 50},
{"key": "variant-2", "rollout_percentage": 50},
]
},
"payloads": {"variant-1": '{"some": "value"}'},
},
}
self.client.feature_flags = [basic_flag]
flag_result = self.client.get_feature_flag_result(
"person-flag", "distinct_id", person_properties={"region": "USA"}
)
self.assertEqual(flag_result.enabled, True)
self.assertEqual(flag_result.variant, "variant-1")
self.assertEqual(flag_result.get_value(), "variant-1")
self.assertEqual(flag_result.payload, {"some": "value"})
patch_capture.assert_called_with(
"$feature_flag_called",
distinct_id="distinct_id",
properties={
"$feature_flag": "person-flag",
"$feature_flag_response": "variant-1",
"locally_evaluated": True,
"$feature/person-flag": "variant-1",
"$feature_flag_payload": {"some": "value"},
},
groups={},
disable_geoip=None,
)
# Verify error property is NOT present on successful evaluation
captured_properties = patch_capture.call_args[1]["properties"]
self.assertNotIn("$feature_flag_error", captured_properties)
another_flag_result = self.client.get_feature_flag_result(
"person-flag", "another-distinct-id", person_properties={"region": "USA"}
)
self.assertEqual(another_flag_result.enabled, True)
self.assertEqual(another_flag_result.variant, "variant-2")
self.assertEqual(another_flag_result.get_value(), "variant-2")
self.assertIsNone(another_flag_result.payload)
patch_capture.assert_called_with(
"$feature_flag_called",
distinct_id="another-distinct-id",
properties={
"$feature_flag": "person-flag",
"$feature_flag_response": "variant-2",
"locally_evaluated": True,
"$feature/person-flag": "variant-2",
},
groups={},
disable_geoip=None,
)
@mock.patch("hanzo_insights.client.flags")
@mock.patch.object(Client, "capture")
def test_get_feature_flag_result_boolean_decide(self, patch_capture, patch_flags):
patch_flags.return_value = {
"flags": {
"person-flag": {
"key": "person-flag",
"enabled": True,
"variant": None,
"reason": {
"description": "Matched condition set 1",
},
"metadata": {
"id": 23,
"version": 42,
"payload": "300",
},
},
},
}
flag_result = self.client.get_feature_flag_result(
"person-flag", "some-distinct-id"
)
self.assertEqual(flag_result.enabled, True)
self.assertEqual(flag_result.variant, None)
self.assertEqual(flag_result.payload, 300)
patch_capture.assert_called_with(
"$feature_flag_called",
distinct_id="some-distinct-id",
properties={
"$feature_flag": "person-flag",
"$feature_flag_response": True,
"locally_evaluated": False,
"$feature/person-flag": True,
"$feature_flag_reason": "Matched condition set 1",
"$feature_flag_id": 23,
"$feature_flag_version": 42,
"$feature_flag_payload": 300,
},
groups={},
disable_geoip=None,
)
# Verify error property is NOT present on successful evaluation
captured_properties = patch_capture.call_args[1]["properties"]
self.assertNotIn("$feature_flag_error", captured_properties)
@mock.patch("hanzo_insights.client.flags")
@mock.patch.object(Client, "capture")
def test_get_feature_flag_result_variant_decide(self, patch_capture, patch_flags):
patch_flags.return_value = {
"flags": {
"person-flag": {
"key": "person-flag",
"enabled": True,
"variant": "variant-1",
"reason": {
"description": "Matched condition set 1",
},
"metadata": {
"id": 1,
"version": 2,
"payload": "[1, 2, 3]",
},
},
},
}
flag_result = self.client.get_feature_flag_result("person-flag", "distinct_id")
self.assertEqual(flag_result.enabled, True)
self.assertEqual(flag_result.variant, "variant-1")
self.assertEqual(flag_result.get_value(), "variant-1")
self.assertEqual(flag_result.payload, [1, 2, 3])
patch_capture.assert_called_with(
"$feature_flag_called",
distinct_id="distinct_id",
properties={
"$feature_flag": "person-flag",
"$feature_flag_response": "variant-1",
"locally_evaluated": False,
"$feature/person-flag": "variant-1",
"$feature_flag_reason": "Matched condition set 1",
"$feature_flag_id": 1,
"$feature_flag_version": 2,
"$feature_flag_payload": [1, 2, 3],
},
groups={},
disable_geoip=None,
)
# Verify error property is NOT present on successful evaluation
captured_properties = patch_capture.call_args[1]["properties"]
self.assertNotIn("$feature_flag_error", captured_properties)
@mock.patch("hanzo_insights.client.flags")
@mock.patch.object(Client, "capture")
def test_get_feature_flag_result_unknown_flag(self, patch_capture, patch_flags):
patch_flags.return_value = {
"flags": {
"person-flag": {
"key": "person-flag",
"enabled": True,
"variant": None,
"reason": {
"description": "Matched condition set 1",
},
"metadata": {
"id": 23,
"version": 42,
"payload": "300",
},
},
},
}
flag_result = self.client.get_feature_flag_result(
"no-person-flag", "some-distinct-id"
)
self.assertIsNone(flag_result)
patch_capture.assert_called_with(
"$feature_flag_called",
distinct_id="some-distinct-id",
properties={
"$feature_flag": "no-person-flag",
"$feature_flag_response": None,
"locally_evaluated": False,
"$feature/no-person-flag": None,
"$feature_flag_error": FeatureFlagError.FLAG_MISSING,
},
groups={},
disable_geoip=None,
)
@mock.patch("hanzo_insights.client.flags")
@mock.patch.object(Client, "capture")
def test_get_feature_flag_result_with_errors_while_computing_flags(
self, patch_capture, patch_flags
):
"""Test that errors_while_computing_flags is included in the $feature_flag_called event.
When the server returns errorsWhileComputingFlags=true, it indicates that there
was an error computing one or more flags. We include this in the event so users
can identify and debug flag evaluation issues.
"""
patch_flags.return_value = {
"flags": {
"my-flag": {
"key": "my-flag",
"enabled": True,
"variant": None,
"reason": {"description": "Matched condition set 1"},
"metadata": {"id": 1, "version": 1, "payload": None},
},
},
"requestId": "test-request-id-789",
"errorsWhileComputingFlags": True,
}
flag_result = self.client.get_feature_flag_result("my-flag", "some-distinct-id")
self.assertEqual(flag_result.enabled, True)
patch_capture.assert_called_with(
"$feature_flag_called",
distinct_id="some-distinct-id",
properties={
"$feature_flag": "my-flag",
"$feature_flag_response": True,
"locally_evaluated": False,
"$feature/my-flag": True,
"$feature_flag_request_id": "test-request-id-789",
"$feature_flag_reason": "Matched condition set 1",
"$feature_flag_id": 1,
"$feature_flag_version": 1,
"$feature_flag_error": FeatureFlagError.ERRORS_WHILE_COMPUTING,
},
groups={},
disable_geoip=None,
)
@mock.patch("hanzo_insights.client.flags")
@mock.patch.object(Client, "capture")
def test_get_feature_flag_result_flag_not_in_response(
self, patch_capture, patch_flags
):
"""Test that when a flag is not in the API response, we capture flag_missing error.
This happens when a flag doesn't exist or the user doesn't match any conditions.
"""
patch_flags.return_value = {
"flags": {
"other-flag": {
"key": "other-flag",
"enabled": True,
"variant": None,
"reason": {"description": "Matched condition set 1"},
"metadata": {"id": 1, "version": 1, "payload": None},
},
},
"requestId": "test-request-id-456",
}
flag_result = self.client.get_feature_flag_result(
"missing-flag", "some-distinct-id"
)
self.assertIsNone(flag_result)
patch_capture.assert_called_with(
"$feature_flag_called",
distinct_id="some-distinct-id",
properties={
"$feature_flag": "missing-flag",
"$feature_flag_response": None,
"locally_evaluated": False,
"$feature/missing-flag": None,
"$feature_flag_request_id": "test-request-id-456",
"$feature_flag_error": FeatureFlagError.FLAG_MISSING,
},
groups={},
disable_geoip=None,
)
@mock.patch("hanzo_insights.client.flags")
@mock.patch.object(Client, "capture")
def test_get_feature_flag_result_errors_computing_and_flag_missing(
self, patch_capture, patch_flags
):
"""Test that both errors are reported when errorsWhileComputingFlags=true AND flag is missing.
This can happen when the server encounters errors computing flags AND the requested
flag is not in the response. Both conditions should be reported for debugging.
"""
patch_flags.return_value = {
"flags": {}, # Flag is missing
"requestId": "test-request-id-999",
"errorsWhileComputingFlags": True, # But errors also occurred
}
flag_result = self.client.get_feature_flag_result(
"missing-flag", "some-distinct-id"
)
self.assertIsNone(flag_result)
patch_capture.assert_called_with(
"$feature_flag_called",
distinct_id="some-distinct-id",
properties={
"$feature_flag": "missing-flag",
"$feature_flag_response": None,
"locally_evaluated": False,
"$feature/missing-flag": None,
"$feature_flag_request_id": "test-request-id-999",
"$feature_flag_error": f"{FeatureFlagError.ERRORS_WHILE_COMPUTING},{FeatureFlagError.FLAG_MISSING}",
},
groups={},
disable_geoip=None,
)
@mock.patch("hanzo_insights.client.flags")
@mock.patch.object(Client, "capture")
def test_get_feature_flag_result_unknown_error(self, patch_capture, patch_flags):
"""Test that unexpected exceptions are captured as unknown_error."""
patch_flags.side_effect = Exception("Unexpected error")
flag_result = self.client.get_feature_flag_result("my-flag", "some-distinct-id")
self.assertIsNone(flag_result)
patch_capture.assert_called_with(
"$feature_flag_called",
distinct_id="some-distinct-id",
properties={
"$feature_flag": "my-flag",
"$feature_flag_response": None,
"locally_evaluated": False,
"$feature/my-flag": None,
"$feature_flag_error": FeatureFlagError.UNKNOWN_ERROR,
},
groups={},
disable_geoip=None,
)
@mock.patch("hanzo_insights.client.flags")
@mock.patch.object(Client, "capture")
def test_get_feature_flag_result_timeout_error(self, patch_capture, patch_flags):
"""Test that timeout errors are captured specifically."""
from hanzo_insights.request import RequestsTimeout
patch_flags.side_effect = RequestsTimeout("Request timed out")
flag_result = self.client.get_feature_flag_result("my-flag", "some-distinct-id")
self.assertIsNone(flag_result)
patch_capture.assert_called_with(
"$feature_flag_called",
distinct_id="some-distinct-id",
properties={
"$feature_flag": "my-flag",
"$feature_flag_response": None,
"locally_evaluated": False,
"$feature/my-flag": None,
"$feature_flag_error": FeatureFlagError.TIMEOUT,
},
groups={},
disable_geoip=None,
)
@mock.patch("hanzo_insights.client.flags")
@mock.patch.object(Client, "capture")
def test_get_feature_flag_result_connection_error(self, patch_capture, patch_flags):
"""Test that connection errors are captured specifically."""
from hanzo_insights.request import RequestsConnectionError
patch_flags.side_effect = RequestsConnectionError("Connection refused")
flag_result = self.client.get_feature_flag_result("my-flag", "some-distinct-id")
self.assertIsNone(flag_result)
patch_capture.assert_called_with(
"$feature_flag_called",
distinct_id="some-distinct-id",
properties={
"$feature_flag": "my-flag",
"$feature_flag_response": None,
"locally_evaluated": False,
"$feature/my-flag": None,
"$feature_flag_error": FeatureFlagError.CONNECTION_ERROR,
},
groups={},
disable_geoip=None,
)
@mock.patch("hanzo_insights.client.flags")
@mock.patch.object(Client, "capture")
def test_get_feature_flag_result_api_error(self, patch_capture, patch_flags):
"""Test that API errors include the status code."""
from hanzo_insights.request import APIError
patch_flags.side_effect = APIError(500, "Internal server error")
flag_result = self.client.get_feature_flag_result("my-flag", "some-distinct-id")
self.assertIsNone(flag_result)
patch_capture.assert_called_with(
"$feature_flag_called",
distinct_id="some-distinct-id",
properties={
"$feature_flag": "my-flag",
"$feature_flag_response": None,
"locally_evaluated": False,
"$feature/my-flag": None,
"$feature_flag_error": FeatureFlagError.api_error(500),
},
groups={},
disable_geoip=None,
)
@mock.patch("hanzo_insights.client.flags")
@mock.patch.object(Client, "capture")
def test_get_feature_flag_result_quota_limited(self, patch_capture, patch_flags):
"""Test that quota limit errors are captured specifically."""
from hanzo_insights.request import QuotaLimitError
patch_flags.side_effect = QuotaLimitError(429, "Rate limit exceeded")
flag_result = self.client.get_feature_flag_result("my-flag", "some-distinct-id")
self.assertIsNone(flag_result)
patch_capture.assert_called_with(
"$feature_flag_called",
distinct_id="some-distinct-id",
properties={
"$feature_flag": "my-flag",
"$feature_flag_response": None,
"locally_evaluated": False,
"$feature/my-flag": None,
"$feature_flag_error": FeatureFlagError.QUOTA_LIMITED,
},
groups={},
disable_geoip=None,
)
class TestFeatureFlagErrorWithStaleCacheFallback(unittest.TestCase):
"""Tests for stale cache fallback behavior when flag evaluation fails.
When the Insights API is unavailable (timeout, connection error, etc.), the SDK
falls back to stale cached flag values if available. These tests verify that:
1. The stale cached value is returned when an error occurs
2. The $feature_flag_error property is still set (for debugging)
3. The response reflects the cached value, not None
"""
def set_fail(self, e, batch):
"""Mark the failure handler"""
self.failed = True
def setUp(self):
self.failed = False
# Create client with memory-based flag cache enabled
self.client = Client(
FAKE_TEST_API_KEY,
on_error=self.set_fail,
flag_fallback_cache_url="memory://local/?ttl=300&size=10000",
)
def _populate_stale_cache(self, distinct_id, flag_key, flag_result):
"""Pre-populate the flag cache with a value that will be used for stale fallback."""
self.client.flag_cache.set_cached_flag(
distinct_id,
flag_key,
flag_result,
flag_definition_version=self.client.flag_definition_version,
)
@mock.patch("hanzo_insights.client.flags")
@mock.patch.object(Client, "capture")
def test_timeout_error_returns_stale_cached_value(self, patch_capture, patch_flags):
"""Test that timeout errors return stale cached value when available."""
from hanzo_insights.request import RequestsTimeout
# Pre-populate cache with a flag result
cached_result = FeatureFlagResult.from_value_and_payload(
"my-flag", "cached-variant", '{"from": "cache"}'
)
self._populate_stale_cache("some-distinct-id", "my-flag", cached_result)
# Simulate timeout error
patch_flags.side_effect = RequestsTimeout("Request timed out")
flag_result = self.client.get_feature_flag_result("my-flag", "some-distinct-id")
# Should return the stale cached value
self.assertIsNotNone(flag_result)
self.assertEqual(flag_result.variant, "cached-variant")
self.assertEqual(flag_result.payload, {"from": "cache"})
# Error should still be tracked for debugging
patch_capture.assert_called_with(
"$feature_flag_called",
distinct_id="some-distinct-id",
properties={
"$feature_flag": "my-flag",
"$feature_flag_response": "cached-variant",
"locally_evaluated": False,
"$feature/my-flag": "cached-variant",
"$feature_flag_payload": {"from": "cache"},
"$feature_flag_error": FeatureFlagError.TIMEOUT,
},
groups={},
disable_geoip=None,
)
@mock.patch("hanzo_insights.client.flags")
@mock.patch.object(Client, "capture")
def test_connection_error_returns_stale_cached_value(
self, patch_capture, patch_flags
):
"""Test that connection errors return stale cached value when available."""
from hanzo_insights.request import RequestsConnectionError
# Pre-populate cache with a boolean flag result
cached_result = FeatureFlagResult.from_value_and_payload("my-flag", True, None)
self._populate_stale_cache("some-distinct-id", "my-flag", cached_result)
# Simulate connection error
patch_flags.side_effect = RequestsConnectionError("Connection refused")
flag_result = self.client.get_feature_flag_result("my-flag", "some-distinct-id")
# Should return the stale cached value
self.assertIsNotNone(flag_result)
self.assertEqual(flag_result.enabled, True)
self.assertIsNone(flag_result.variant)
# Error should still be tracked
patch_capture.assert_called_with(
"$feature_flag_called",
distinct_id="some-distinct-id",
properties={
"$feature_flag": "my-flag",
"$feature_flag_response": True,
"locally_evaluated": False,
"$feature/my-flag": True,
"$feature_flag_error": FeatureFlagError.CONNECTION_ERROR,
},
groups={},
disable_geoip=None,
)
@mock.patch("hanzo_insights.client.flags")
@mock.patch.object(Client, "capture")
def test_api_error_returns_stale_cached_value(self, patch_capture, patch_flags):
"""Test that API errors return stale cached value when available."""
from hanzo_insights.request import APIError
# Pre-populate cache
cached_result = FeatureFlagResult.from_value_and_payload(
"my-flag", "control", None
)
self._populate_stale_cache("some-distinct-id", "my-flag", cached_result)
# Simulate API error
patch_flags.side_effect = APIError(503, "Service unavailable")
flag_result = self.client.get_feature_flag_result("my-flag", "some-distinct-id")
# Should return the stale cached value
self.assertIsNotNone(flag_result)
self.assertEqual(flag_result.variant, "control")
# Error should still be tracked with status code
patch_capture.assert_called_with(
"$feature_flag_called",
distinct_id="some-distinct-id",
properties={
"$feature_flag": "my-flag",
"$feature_flag_response": "control",
"locally_evaluated": False,
"$feature/my-flag": "control",
"$feature_flag_error": FeatureFlagError.api_error(503),
},
groups={},
disable_geoip=None,
)
@mock.patch("hanzo_insights.client.flags")
@mock.patch.object(Client, "capture")
def test_error_without_cache_returns_none(self, patch_capture, patch_flags):
"""Test that errors return None when no stale cache is available."""
from hanzo_insights.request import RequestsTimeout
# Do NOT populate cache - no fallback available
patch_flags.side_effect = RequestsTimeout("Request timed out")
flag_result = self.client.get_feature_flag_result("my-flag", "some-distinct-id")
# Should return None since no cache available
self.assertIsNone(flag_result)
# Error should still be tracked
patch_capture.assert_called_with(
"$feature_flag_called",
distinct_id="some-distinct-id",
properties={
"$feature_flag": "my-flag",
"$feature_flag_response": None,
"locally_evaluated": False,
"$feature/my-flag": None,
"$feature_flag_error": FeatureFlagError.TIMEOUT,
},
groups={},
disable_geoip=None,
)
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,612 @@
"""
Tests for FlagDefinitionCacheProvider functionality.
These tests follow the patterns from the TypeScript implementation in insights-js/packages/node.
"""
import threading
import unittest
from typing import Optional
from unittest import mock
from hanzo_insights.client import Client
from hanzo_insights.flag_definition_cache import (
FlagDefinitionCacheData,
FlagDefinitionCacheProvider,
)
from hanzo_insights.request import GetResponse
from hanzo_insights.test.test_utils import FAKE_TEST_API_KEY
class MockCacheProvider:
"""A mock implementation of FlagDefinitionCacheProvider for testing."""
def __init__(self):
self.stored_data: Optional[FlagDefinitionCacheData] = None
self.should_fetch_return_value = True
self.get_call_count = 0
self.should_fetch_call_count = 0
self.on_received_call_count = 0
self.shutdown_call_count = 0
self.should_fetch_error: Optional[Exception] = None
self.get_error: Optional[Exception] = None
self.on_received_error: Optional[Exception] = None
self.shutdown_error: Optional[Exception] = None
def get_flag_definitions(self) -> Optional[FlagDefinitionCacheData]:
self.get_call_count += 1
if self.get_error:
raise self.get_error
return self.stored_data
def should_fetch_flag_definitions(self) -> bool:
self.should_fetch_call_count += 1
if self.should_fetch_error:
raise self.should_fetch_error
return self.should_fetch_return_value
def on_flag_definitions_received(self, data: FlagDefinitionCacheData) -> None:
self.on_received_call_count += 1
if self.on_received_error:
raise self.on_received_error
self.stored_data = data
def shutdown(self) -> None:
self.shutdown_call_count += 1
if self.shutdown_error:
raise self.shutdown_error
class TestFlagDefinitionCacheProvider(unittest.TestCase):
"""Tests for the FlagDefinitionCacheProvider protocol."""
@classmethod
def setUpClass(cls):
# Prevent real HTTP requests
cls.client_post_patcher = mock.patch("hanzo_insights.client.batch_post")
cls.consumer_post_patcher = mock.patch("hanzo_insights.consumer.batch_post")
cls.client_post_patcher.start()
cls.consumer_post_patcher.start()
@classmethod
def tearDownClass(cls):
cls.client_post_patcher.stop()
cls.consumer_post_patcher.stop()
def setUp(self):
self.cache_provider = MockCacheProvider()
self.sample_flags_data: FlagDefinitionCacheData = {
"flags": [
{"key": "test-flag", "active": True, "filters": {}},
{"key": "another-flag", "active": False, "filters": {}},
],
"group_type_mapping": {"0": "company", "1": "project"},
"cohorts": {"1": {"properties": []}},
}
def tearDown(self):
# Ensure client cleanup
pass
def _create_client_with_cache(self) -> Client:
"""Create a client with the mock cache provider."""
return Client(
FAKE_TEST_API_KEY,
personal_api_key="test-personal-key",
flag_definition_cache_provider=self.cache_provider,
sync_mode=True,
enable_local_evaluation=False, # Disable poller for tests
)
class TestCacheInitialization(TestFlagDefinitionCacheProvider):
"""Tests for cache initialization behavior."""
@mock.patch("hanzo_insights.client.get")
def test_uses_cached_data_when_should_fetch_returns_false(self, mock_get):
"""When should_fetch returns False and cache has data, use cached data."""
self.cache_provider.should_fetch_return_value = False
self.cache_provider.stored_data = self.sample_flags_data
client = self._create_client_with_cache()
client._load_feature_flags()
# Should not call API
mock_get.assert_not_called()
# Should have called cache methods
self.assertEqual(self.cache_provider.should_fetch_call_count, 1)
self.assertEqual(self.cache_provider.get_call_count, 1)
# Flags should be loaded from cache
self.assertEqual(len(client.feature_flags), 2)
self.assertEqual(client.feature_flags[0]["key"], "test-flag")
client.join()
@mock.patch("hanzo_insights.client.get")
def test_fetches_from_api_when_should_fetch_returns_true(self, mock_get):
"""When should_fetch returns True, fetch from API."""
self.cache_provider.should_fetch_return_value = True
mock_get.return_value = GetResponse(
data=self.sample_flags_data, etag="test-etag", not_modified=False
)
client = self._create_client_with_cache()
client._load_feature_flags()
# Should call API
mock_get.assert_called_once()
# Should have called should_fetch but not get
self.assertEqual(self.cache_provider.should_fetch_call_count, 1)
self.assertEqual(self.cache_provider.get_call_count, 0)
# Should have called on_received to store in cache
self.assertEqual(self.cache_provider.on_received_call_count, 1)
client.join()
@mock.patch("hanzo_insights.client.get")
def test_emergency_fallback_when_cache_empty_and_no_flags(self, mock_get):
"""When should_fetch=False but cache is empty and no flags loaded, fetch anyway."""
self.cache_provider.should_fetch_return_value = False
self.cache_provider.stored_data = None # Empty cache
mock_get.return_value = GetResponse(
data=self.sample_flags_data, etag="test-etag", not_modified=False
)
client = self._create_client_with_cache()
client._load_feature_flags()
# Should call API due to emergency fallback
mock_get.assert_called_once()
# Should have called on_received
self.assertEqual(self.cache_provider.on_received_call_count, 1)
client.join()
@mock.patch("hanzo_insights.client.get")
def test_preserves_existing_flags_when_cache_returns_none(self, mock_get):
"""When cache returns None but client has flags, preserve existing flags."""
self.cache_provider.should_fetch_return_value = False
self.cache_provider.stored_data = None # Empty cache
client = self._create_client_with_cache()
# Pre-load flags (simulating a previous successful fetch)
client.feature_flags = self.sample_flags_data["flags"]
client.group_type_mapping = self.sample_flags_data["group_type_mapping"]
client.cohorts = self.sample_flags_data["cohorts"]
client._load_feature_flags()
# Should NOT call API since we already have flags
mock_get.assert_not_called()
# Existing flags should be preserved
self.assertEqual(len(client.feature_flags), 2)
self.assertEqual(client.feature_flags[0]["key"], "test-flag")
client.join()
class TestFetchCoordination(TestFlagDefinitionCacheProvider):
"""Tests for fetch coordination between workers."""
@mock.patch("hanzo_insights.client.get")
def test_calls_should_fetch_before_each_poll(self, mock_get):
"""should_fetch_flag_definitions is called before each poll cycle."""
self.cache_provider.should_fetch_return_value = True
mock_get.return_value = GetResponse(
data=self.sample_flags_data, etag="test-etag", not_modified=False
)
client = self._create_client_with_cache()
# First poll
client._load_feature_flags()
self.assertEqual(self.cache_provider.should_fetch_call_count, 1)
# Second poll
client._load_feature_flags()
self.assertEqual(self.cache_provider.should_fetch_call_count, 2)
client.join()
@mock.patch("hanzo_insights.client.get")
def test_does_not_call_on_received_when_fetch_skipped(self, mock_get):
"""on_flag_definitions_received is NOT called when fetch is skipped."""
self.cache_provider.should_fetch_return_value = False
self.cache_provider.stored_data = self.sample_flags_data
client = self._create_client_with_cache()
client._load_feature_flags()
# Should not call on_received since we didn't fetch
self.assertEqual(self.cache_provider.on_received_call_count, 0)
client.join()
@mock.patch("hanzo_insights.client.get")
def test_stores_data_in_cache_after_api_fetch(self, mock_get):
"""on_flag_definitions_received receives the fetched data."""
self.cache_provider.should_fetch_return_value = True
mock_get.return_value = GetResponse(
data=self.sample_flags_data, etag="test-etag", not_modified=False
)
client = self._create_client_with_cache()
client._load_feature_flags()
# Should have stored data in cache
self.assertEqual(self.cache_provider.on_received_call_count, 1)
self.assertIsNotNone(self.cache_provider.stored_data)
self.assertEqual(len(self.cache_provider.stored_data["flags"]), 2)
client.join()
@mock.patch("hanzo_insights.client.get")
def test_304_not_modified_does_not_update_cache(self, mock_get):
"""When API returns 304 Not Modified, cache should not be updated."""
self.cache_provider.should_fetch_return_value = True
# First fetch to populate flags and ETag
mock_get.return_value = GetResponse(
data=self.sample_flags_data, etag="test-etag", not_modified=False
)
client = self._create_client_with_cache()
client._load_feature_flags()
# Verify initial fetch worked
self.assertEqual(self.cache_provider.on_received_call_count, 1)
self.assertEqual(len(client.feature_flags), 2)
# Second fetch returns 304 Not Modified
mock_get.return_value = GetResponse(
data=None, etag="test-etag", not_modified=True
)
client._load_feature_flags()
# API was called twice
self.assertEqual(mock_get.call_count, 2)
# should_fetch was called twice
self.assertEqual(self.cache_provider.should_fetch_call_count, 2)
# on_received should NOT be called again (304 = no new data)
self.assertEqual(self.cache_provider.on_received_call_count, 1)
# Flags should still be present
self.assertEqual(len(client.feature_flags), 2)
client.join()
class TestErrorHandling(TestFlagDefinitionCacheProvider):
"""Tests for error handling in cache provider operations."""
@mock.patch("hanzo_insights.client.get")
def test_should_fetch_error_defaults_to_fetching(self, mock_get):
"""When should_fetch throws an error, default to fetching from API."""
self.cache_provider.should_fetch_error = Exception("Lock acquisition failed")
mock_get.return_value = GetResponse(
data=self.sample_flags_data, etag="test-etag", not_modified=False
)
client = self._create_client_with_cache()
client._load_feature_flags()
# Should still fetch from API
mock_get.assert_called_once()
# Flags should be loaded
self.assertEqual(len(client.feature_flags), 2)
client.join()
@mock.patch("hanzo_insights.client.get")
def test_get_error_falls_back_to_api_fetch(self, mock_get):
"""When get_flag_definitions throws an error, fetch from API."""
self.cache_provider.should_fetch_return_value = False
self.cache_provider.get_error = Exception("Cache read failed")
mock_get.return_value = GetResponse(
data=self.sample_flags_data, etag="test-etag", not_modified=False
)
client = self._create_client_with_cache()
client._load_feature_flags()
# Should fall back to API
mock_get.assert_called_once()
client.join()
@mock.patch("hanzo_insights.client.get")
def test_on_received_error_keeps_flags_in_memory(self, mock_get):
"""When on_flag_definitions_received throws, flags are still in memory."""
self.cache_provider.should_fetch_return_value = True
self.cache_provider.on_received_error = Exception("Cache write failed")
mock_get.return_value = GetResponse(
data=self.sample_flags_data, etag="test-etag", not_modified=False
)
client = self._create_client_with_cache()
client._load_feature_flags()
# Flags should still be loaded in memory despite cache error
self.assertEqual(len(client.feature_flags), 2)
self.assertEqual(client.feature_flags[0]["key"], "test-flag")
client.join()
@mock.patch("hanzo_insights.client.get")
def test_shutdown_error_is_logged_but_continues(self, mock_get):
"""When shutdown throws an error, it's logged but shutdown continues."""
self.cache_provider.shutdown_error = Exception("Lock release failed")
mock_get.return_value = GetResponse(
data=self.sample_flags_data, etag="test-etag", not_modified=False
)
client = self._create_client_with_cache()
client._load_feature_flags()
# Should not raise when joining
client.join()
# Shutdown was called
self.assertEqual(self.cache_provider.shutdown_call_count, 1)
class TestShutdownLifecycle(TestFlagDefinitionCacheProvider):
"""Tests for shutdown lifecycle."""
@mock.patch("hanzo_insights.client.get")
def test_shutdown_calls_cache_provider_shutdown(self, mock_get):
"""Client shutdown calls cache provider shutdown."""
mock_get.return_value = GetResponse(
data=self.sample_flags_data, etag="test-etag", not_modified=False
)
client = self._create_client_with_cache()
client._load_feature_flags()
# Shutdown
client.join()
self.assertEqual(self.cache_provider.shutdown_call_count, 1)
@mock.patch("hanzo_insights.client.get")
def test_shutdown_called_even_without_fetching(self, mock_get):
"""Shutdown is called even when cache was used instead of fetching."""
self.cache_provider.should_fetch_return_value = False
self.cache_provider.stored_data = self.sample_flags_data
client = self._create_client_with_cache()
client._load_feature_flags()
client.join()
# Shutdown should still be called
self.assertEqual(self.cache_provider.shutdown_call_count, 1)
@mock.patch("hanzo_insights.client.get")
def test_multiple_join_calls_only_shutdown_once(self, mock_get):
"""Calling join() multiple times should only call cache provider shutdown once."""
mock_get.return_value = GetResponse(
data=self.sample_flags_data, etag="test-etag", not_modified=False
)
client = self._create_client_with_cache()
client._load_feature_flags()
# Call join multiple times
client.join()
client.join()
client.join()
# Shutdown should be called each time (current behavior - no guard)
# This test documents the current behavior
self.assertGreaterEqual(self.cache_provider.shutdown_call_count, 1)
class TestBackwardCompatibility(TestFlagDefinitionCacheProvider):
"""Tests for backward compatibility without cache provider."""
@mock.patch("hanzo_insights.client.get")
def test_works_without_cache_provider(self, mock_get):
"""Client works normally without a cache provider configured."""
mock_get.return_value = GetResponse(
data=self.sample_flags_data, etag="test-etag", not_modified=False
)
# Create client without cache provider
client = Client(
FAKE_TEST_API_KEY,
personal_api_key="test-personal-key",
sync_mode=True,
enable_local_evaluation=False,
)
client._load_feature_flags()
# Should fetch from API
mock_get.assert_called_once()
# Flags should be loaded
self.assertEqual(len(client.feature_flags), 2)
client.join()
class TestDataIntegrity(TestFlagDefinitionCacheProvider):
"""Tests for data integrity between cache and client state."""
@mock.patch("hanzo_insights.client.get")
def test_cached_flags_available_for_evaluation(self, mock_get):
"""Flags loaded from cache are available for local evaluation."""
self.cache_provider.should_fetch_return_value = False
self.cache_provider.stored_data = {
"flags": [
{
"key": "test-flag",
"active": True,
"filters": {
"groups": [
{
"properties": [],
"rollout_percentage": 100,
}
]
},
}
],
"group_type_mapping": {},
"cohorts": {},
}
client = self._create_client_with_cache()
client._load_feature_flags()
# Flag should be accessible
self.assertEqual(len(client.feature_flags), 1)
self.assertEqual(client.feature_flags_by_key["test-flag"]["key"], "test-flag")
client.join()
@mock.patch("hanzo_insights.client.get")
def test_group_type_mapping_loaded_from_cache(self, mock_get):
"""Group type mapping is correctly loaded from cache."""
self.cache_provider.should_fetch_return_value = False
self.cache_provider.stored_data = self.sample_flags_data
client = self._create_client_with_cache()
client._load_feature_flags()
self.assertEqual(client.group_type_mapping["0"], "company")
self.assertEqual(client.group_type_mapping["1"], "project")
client.join()
@mock.patch("hanzo_insights.client.get")
def test_cohorts_loaded_from_cache(self, mock_get):
"""Cohorts are correctly loaded from cache."""
self.cache_provider.should_fetch_return_value = False
self.cache_provider.stored_data = self.sample_flags_data
client = self._create_client_with_cache()
client._load_feature_flags()
self.assertIn("1", client.cohorts)
client.join()
@mock.patch("hanzo_insights.client.get")
def test_cache_updated_when_api_returns_new_data(self, mock_get):
"""State transition: cache has old data -> API returns new -> cache updated."""
# Start with old cached data
old_flags_data: FlagDefinitionCacheData = {
"flags": [{"key": "old-flag", "active": True, "filters": {}}],
"group_type_mapping": {},
"cohorts": {},
}
self.cache_provider.stored_data = old_flags_data
self.cache_provider.should_fetch_return_value = False
client = self._create_client_with_cache()
# First load from cache
client._load_feature_flags()
self.assertEqual(client.feature_flags[0]["key"], "old-flag")
self.assertEqual(self.cache_provider.on_received_call_count, 0)
# Now trigger API fetch with new data
self.cache_provider.should_fetch_return_value = True
new_flags_data: FlagDefinitionCacheData = {
"flags": [{"key": "new-flag", "active": True, "filters": {}}],
"group_type_mapping": {"0": "company"},
"cohorts": {"1": {"properties": []}},
}
mock_get.return_value = GetResponse(
data=new_flags_data, etag="new-etag", not_modified=False
)
client._load_feature_flags()
# Verify new flags loaded
self.assertEqual(client.feature_flags[0]["key"], "new-flag")
self.assertEqual(client.group_type_mapping["0"], "company")
# Verify cache was updated
self.assertEqual(self.cache_provider.on_received_call_count, 1)
self.assertEqual(self.cache_provider.stored_data["flags"][0]["key"], "new-flag")
client.join()
class TestConcurrency(TestFlagDefinitionCacheProvider):
"""Tests for thread safety and concurrent access."""
@mock.patch("hanzo_insights.client.get")
def test_concurrent_load_feature_flags_is_thread_safe(self, mock_get):
"""Multiple threads calling _load_feature_flags should not cause errors."""
mock_get.return_value = GetResponse(
data=self.sample_flags_data, etag="test-etag", not_modified=False
)
client = self._create_client_with_cache()
errors = []
def load_flags():
try:
client._load_feature_flags()
except Exception as e:
errors.append(e)
# Launch 5 threads concurrently
threads = [threading.Thread(target=load_flags) for _ in range(5)]
for t in threads:
t.start()
for t in threads:
t.join()
# Should complete without errors
self.assertEqual(len(errors), 0, f"Unexpected errors: {errors}")
# Flags should be loaded
self.assertIsNotNone(client.feature_flags)
self.assertEqual(len(client.feature_flags), 2)
client.join()
class TestProtocolCompliance(unittest.TestCase):
"""Tests for Protocol compliance."""
def test_mock_provider_is_protocol_instance(self):
"""MockCacheProvider satisfies FlagDefinitionCacheProvider protocol."""
provider = MockCacheProvider()
self.assertIsInstance(provider, FlagDefinitionCacheProvider)
def test_incomplete_provider_is_not_protocol_instance(self):
"""Class missing methods is not a FlagDefinitionCacheProvider."""
class IncompleteProvider:
def get_flag_definitions(self):
return None
provider = IncompleteProvider()
self.assertNotIsInstance(provider, FlagDefinitionCacheProvider)
if __name__ == "__main__":
unittest.main()
@@ -1,10 +1,10 @@
import unittest
from posthog import Posthog
from hanzo_insights import Insights
class TestModule(unittest.TestCase):
posthog = None
client = None
def _assert_enqueue_result(self, result):
self.assertEqual(type(result[0]), str)
@@ -14,19 +14,19 @@ class TestModule(unittest.TestCase):
def setUp(self):
self.failed = False
self.posthog = Posthog(
self.client = Insights(
"testsecret", host="http://localhost:8000", on_error=self.failed
)
def test_track(self):
res = self.posthog.capture("python module event", distinct_id="distinct_id")
res = self.client.capture("python module event", distinct_id="distinct_id")
self._assert_enqueue_result(res)
self.posthog.flush()
self.client.flush()
def test_alias(self):
res = self.posthog.alias("previousId", "distinct_id")
res = self.client.alias("previousId", "distinct_id")
self._assert_enqueue_result(res)
self.posthog.flush()
self.client.flush()
def test_flush(self):
self.posthog.flush()
self.client.flush()
+666
View File
@@ -0,0 +1,666 @@
import json
import unittest
from datetime import date, datetime
import mock
import pytest
import requests
import hanzo_insights.request as request_module
from hanzo_insights.request import (
APIError,
DatetimeSerializer,
GetResponse,
KEEP_ALIVE_SOCKET_OPTIONS,
QuotaLimitError,
_mask_tokens_in_url,
batch_post,
decide,
determine_server_host,
disable_connection_reuse,
enable_keep_alive,
flags,
get,
set_socket_options,
)
from hanzo_insights.test.test_utils import TEST_API_KEY
@pytest.mark.parametrize(
"url, expected",
[
# Token with params after - masks keeping first 10 chars
(
"https://example.com/api/flags?token=phc_abc123xyz789&send_cohorts",
"https://example.com/api/flags?token=phc_abc123...&send_cohorts",
),
# Token at end of URL
(
"https://example.com/api/flags?token=phc_abc123xyz789",
"https://example.com/api/flags?token=phc_abc123...",
),
# No token - unchanged
(
"https://example.com/api/flags?other=value",
"https://example.com/api/flags?other=value",
),
# Short token (<10 chars) - unchanged
(
"https://example.com/api/flags?token=short",
"https://example.com/api/flags?token=short",
),
# Exactly 10 char token - gets ellipsis
(
"https://example.com/api/flags?token=1234567890",
"https://example.com/api/flags?token=1234567890...",
),
],
)
def test_mask_tokens_in_url(url, expected):
assert _mask_tokens_in_url(url) == expected
class TestRequests(unittest.TestCase):
def test_valid_request(self):
res = batch_post(
TEST_API_KEY,
batch=[
{"distinct_id": "distinct_id", "event": "python event", "type": "track"}
],
)
self.assertEqual(res.status_code, 200)
def test_invalid_request_error(self):
self.assertRaises(
Exception, batch_post, "testsecret", "https://t.posthog.com", False, "[{]"
)
def test_invalid_host(self):
self.assertRaises(
Exception, batch_post, "testsecret", "t.posthog.com/", batch=[]
)
def test_datetime_serialization(self):
data = {"created": datetime(2012, 3, 4, 5, 6, 7, 891011)}
result = json.dumps(data, cls=DatetimeSerializer)
self.assertEqual(result, '{"created": "2012-03-04T05:06:07.891011"}')
def test_date_serialization(self):
today = date.today()
data = {"created": today}
result = json.dumps(data, cls=DatetimeSerializer)
expected = '{"created": "%s"}' % today.isoformat()
self.assertEqual(result, expected)
def test_should_not_timeout(self):
res = batch_post(
TEST_API_KEY,
batch=[
{"distinct_id": "distinct_id", "event": "python event", "type": "track"}
],
timeout=15,
)
self.assertEqual(res.status_code, 200)
def test_should_timeout(self):
with self.assertRaises(requests.ReadTimeout):
batch_post(
"key",
batch=[
{
"distinct_id": "distinct_id",
"event": "python event",
"type": "track",
}
],
timeout=0.0001,
)
def test_quota_limited_response(self):
mock_response = requests.Response()
mock_response.status_code = 200
mock_response._content = json.dumps(
{
"quotaLimited": ["feature_flags"],
"featureFlags": {},
"featureFlagPayloads": {},
"errorsWhileComputingFlags": False,
}
).encode("utf-8")
with mock.patch("hanzo_insights.request._session.post", return_value=mock_response):
with self.assertRaises(QuotaLimitError) as cm:
decide("fake_key", "fake_host")
self.assertEqual(cm.exception.status, 200)
self.assertEqual(cm.exception.message, "Feature flags quota limited")
def test_normal_decide_response(self):
mock_response = requests.Response()
mock_response.status_code = 200
mock_response._content = json.dumps(
{
"featureFlags": {"flag1": True},
"featureFlagPayloads": {},
"errorsWhileComputingFlags": False,
}
).encode("utf-8")
with mock.patch("hanzo_insights.request._session.post", return_value=mock_response):
response = decide("fake_key", "fake_host")
self.assertEqual(response["featureFlags"], {"flag1": True})
class TestGet(unittest.TestCase):
"""Unit tests for the get() function HTTP-level behavior."""
@mock.patch("hanzo_insights.request._session.get")
def test_get_returns_data_and_etag(self, mock_get):
"""Test that get() returns GetResponse with data and etag from headers."""
mock_response = requests.Response()
mock_response.status_code = 200
mock_response.headers["ETag"] = '"abc123"'
mock_response._content = json.dumps({"flags": [{"key": "test-flag"}]}).encode(
"utf-8"
)
mock_get.return_value = mock_response
response = get("api_key", "/test-url", host="https://example.com")
self.assertIsInstance(response, GetResponse)
self.assertEqual(response.data, {"flags": [{"key": "test-flag"}]})
self.assertEqual(response.etag, '"abc123"')
self.assertFalse(response.not_modified)
@mock.patch("hanzo_insights.request._session.get")
def test_get_sends_if_none_match_header_when_etag_provided(self, mock_get):
"""Test that If-None-Match header is sent when etag parameter is provided."""
mock_response = requests.Response()
mock_response.status_code = 200
mock_response.headers["ETag"] = '"new-etag"'
mock_response._content = json.dumps({"flags": []}).encode("utf-8")
mock_get.return_value = mock_response
get("api_key", "/test-url", host="https://example.com", etag='"previous-etag"')
call_kwargs = mock_get.call_args[1]
self.assertEqual(call_kwargs["headers"]["If-None-Match"], '"previous-etag"')
@mock.patch("hanzo_insights.request._session.get")
def test_get_does_not_send_if_none_match_when_no_etag(self, mock_get):
"""Test that If-None-Match header is not sent when no etag provided."""
mock_response = requests.Response()
mock_response.status_code = 200
mock_response._content = json.dumps({"flags": []}).encode("utf-8")
mock_get.return_value = mock_response
get("api_key", "/test-url", host="https://example.com")
call_kwargs = mock_get.call_args[1]
self.assertNotIn("If-None-Match", call_kwargs["headers"])
@mock.patch("hanzo_insights.request._session.get")
def test_get_handles_304_not_modified(self, mock_get):
"""Test that 304 Not Modified response returns not_modified=True with no data."""
mock_response = requests.Response()
mock_response.status_code = 304
mock_response.headers["ETag"] = '"unchanged-etag"'
mock_get.return_value = mock_response
response = get(
"api_key", "/test-url", host="https://example.com", etag='"unchanged-etag"'
)
self.assertIsInstance(response, GetResponse)
self.assertIsNone(response.data)
self.assertEqual(response.etag, '"unchanged-etag"')
self.assertTrue(response.not_modified)
@mock.patch("hanzo_insights.request._session.get")
def test_get_304_without_etag_header_uses_request_etag(self, mock_get):
"""Test that 304 response without ETag header falls back to request etag."""
mock_response = requests.Response()
mock_response.status_code = 304
# Server doesn't return ETag header on 304
mock_get.return_value = mock_response
response = get(
"api_key", "/test-url", host="https://example.com", etag='"original-etag"'
)
self.assertTrue(response.not_modified)
self.assertEqual(response.etag, '"original-etag"')
@mock.patch("hanzo_insights.request._session.get")
def test_get_200_without_etag_header(self, mock_get):
"""Test that 200 response without ETag header returns None for etag."""
mock_response = requests.Response()
mock_response.status_code = 200
mock_response._content = json.dumps({"flags": []}).encode("utf-8")
# No ETag header
mock_get.return_value = mock_response
response = get("api_key", "/test-url", host="https://example.com")
self.assertFalse(response.not_modified)
self.assertIsNone(response.etag)
self.assertEqual(response.data, {"flags": []})
@mock.patch("hanzo_insights.request._session.get")
def test_get_error_response_raises_api_error(self, mock_get):
"""Test that error responses raise APIError."""
mock_response = requests.Response()
mock_response.status_code = 401
mock_response._content = json.dumps({"detail": "Unauthorized"}).encode("utf-8")
mock_get.return_value = mock_response
with self.assertRaises(APIError) as ctx:
get("bad_key", "/test-url", host="https://example.com")
self.assertEqual(ctx.exception.status, 401)
self.assertEqual(ctx.exception.message, "Unauthorized")
@mock.patch("hanzo_insights.request._session.get")
def test_get_sends_authorization_header(self, mock_get):
"""Test that Authorization header is sent with Bearer token."""
mock_response = requests.Response()
mock_response.status_code = 200
mock_response._content = json.dumps({}).encode("utf-8")
mock_get.return_value = mock_response
get("my-api-key", "/test-url", host="https://example.com")
call_kwargs = mock_get.call_args[1]
self.assertEqual(call_kwargs["headers"]["Authorization"], "Bearer my-api-key")
@mock.patch("hanzo_insights.request._session.get")
def test_get_sends_user_agent_header(self, mock_get):
"""Test that User-Agent header is sent."""
mock_response = requests.Response()
mock_response.status_code = 200
mock_response._content = json.dumps({}).encode("utf-8")
mock_get.return_value = mock_response
get("api_key", "/test-url", host="https://example.com")
call_kwargs = mock_get.call_args[1]
self.assertIn("User-Agent", call_kwargs["headers"])
self.assertTrue(
call_kwargs["headers"]["User-Agent"].startswith("hanzo-insights-python/")
)
@mock.patch("hanzo_insights.request._session.get")
def test_get_passes_timeout(self, mock_get):
"""Test that timeout parameter is passed to the request."""
mock_response = requests.Response()
mock_response.status_code = 200
mock_response._content = json.dumps({}).encode("utf-8")
mock_get.return_value = mock_response
get("api_key", "/test-url", host="https://example.com", timeout=30)
call_kwargs = mock_get.call_args[1]
self.assertEqual(call_kwargs["timeout"], 30)
@mock.patch("hanzo_insights.request._session.get")
def test_get_constructs_full_url(self, mock_get):
"""Test that host and url are combined correctly."""
mock_response = requests.Response()
mock_response.status_code = 200
mock_response._content = json.dumps({}).encode("utf-8")
mock_get.return_value = mock_response
get("api_key", "/api/flags", host="https://example.com")
call_args = mock_get.call_args[0]
self.assertEqual(call_args[0], "https://example.com/api/flags")
@mock.patch("hanzo_insights.request._session.get")
def test_get_removes_trailing_slash_from_host(self, mock_get):
"""Test that trailing slash is removed from host."""
mock_response = requests.Response()
mock_response.status_code = 200
mock_response._content = json.dumps({}).encode("utf-8")
mock_get.return_value = mock_response
get("api_key", "/api/flags", host="https://example.com/")
call_args = mock_get.call_args[0]
self.assertEqual(call_args[0], "https://example.com/api/flags")
@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.insights.hanzo.ai"),
("https://eu.posthog.com", "https://eu.i.insights.hanzo.ai"),
("https://us.posthog.com", "https://us.i.insights.hanzo.ai"),
("https://app.posthog.com/", "https://us.i.insights.hanzo.ai"),
("https://eu.posthog.com/", "https://eu.i.insights.hanzo.ai"),
("https://us.posthog.com/", "https://us.i.insights.hanzo.ai"),
(None, "https://us.i.insights.hanzo.ai"),
],
)
def test_routing_to_custom_host(host, expected):
assert determine_server_host(host) == expected
def test_enable_keep_alive_sets_socket_options():
try:
enable_keep_alive()
from hanzo_insights.request import _session
adapter = _session.get_adapter("https://example.com")
assert adapter.socket_options == KEEP_ALIVE_SOCKET_OPTIONS
finally:
set_socket_options(None)
def test_set_socket_options_clears_with_none():
try:
enable_keep_alive()
set_socket_options(None)
from hanzo_insights.request import _session
adapter = _session.get_adapter("https://example.com")
assert adapter.socket_options is None
finally:
set_socket_options(None)
def test_disable_connection_reuse_creates_fresh_sessions():
try:
disable_connection_reuse()
session1 = request_module._get_session()
session2 = request_module._get_session()
assert session1 is not session2
finally:
request_module._pooling_enabled = True
def test_set_socket_options_is_idempotent():
try:
enable_keep_alive()
session1 = request_module._session
enable_keep_alive()
session2 = request_module._session
assert session1 is session2
finally:
set_socket_options(None)
class TestFlagsSession(unittest.TestCase):
"""Tests for flags session configuration."""
def test_retry_status_forcelist_excludes_rate_limits(self):
"""Verify 429 (rate limit) is NOT retried - need to wait, not hammer."""
from hanzo_insights.request import RETRY_STATUS_FORCELIST
self.assertNotIn(429, RETRY_STATUS_FORCELIST)
def test_retry_status_forcelist_excludes_quota_errors(self):
"""Verify 402 (payment required/quota) is NOT retried - won't resolve."""
from hanzo_insights.request import RETRY_STATUS_FORCELIST
self.assertNotIn(402, RETRY_STATUS_FORCELIST)
@mock.patch("hanzo_insights.request._get_flags_session")
def test_flags_uses_flags_session(self, mock_get_flags_session):
"""flags() uses the dedicated flags session, not the general session."""
mock_response = requests.Response()
mock_response.status_code = 200
mock_response._content = json.dumps(
{
"featureFlags": {"test-flag": True},
"featureFlagPayloads": {},
"errorsWhileComputingFlags": False,
}
).encode("utf-8")
mock_session = mock.MagicMock()
mock_session.post.return_value = mock_response
mock_get_flags_session.return_value = mock_session
result = flags("test-key", "https://test.posthog.com", distinct_id="user123")
self.assertEqual(result["featureFlags"]["test-flag"], True)
mock_get_flags_session.assert_called_once()
mock_session.post.assert_called_once()
@mock.patch("hanzo_insights.request._get_flags_session")
def test_flags_no_retry_on_quota_limit(self, mock_get_flags_session):
"""flags() raises QuotaLimitError without retrying (at application level)."""
mock_response = requests.Response()
mock_response.status_code = 200
mock_response._content = json.dumps(
{
"quotaLimited": ["feature_flags"],
"featureFlags": {},
"featureFlagPayloads": {},
"errorsWhileComputingFlags": False,
}
).encode("utf-8")
mock_session = mock.MagicMock()
mock_session.post.return_value = mock_response
mock_get_flags_session.return_value = mock_session
with self.assertRaises(QuotaLimitError):
flags("test-key", "https://test.posthog.com", distinct_id="user123")
# QuotaLimitError is raised after response is received, not retried
self.assertEqual(mock_session.post.call_count, 1)
class TestFlagsSessionNetworkRetries(unittest.TestCase):
"""Tests for network failure retries in the flags session."""
def test_flags_session_retry_config_includes_connection_errors(self):
"""
Verify that the flags session is configured to retry on connection errors.
The urllib3 Retry adapter with connect=2 and read=2 automatically
retries on network-level failures (DNS failures, connection refused,
connection reset, etc.) up to 2 times each.
"""
from hanzo_insights.request import _build_flags_session
session = _build_flags_session()
# Get the adapter for https://
adapter = session.get_adapter("https://test.posthog.com")
# Verify retry configuration
retry = adapter.max_retries
self.assertEqual(retry.total, 2, "Should have 2 total retries")
self.assertEqual(retry.connect, 2, "Should retry connection errors twice")
self.assertEqual(retry.read, 2, "Should retry read errors twice")
self.assertIn("POST", retry.allowed_methods, "Should allow POST retries")
def test_flags_session_retries_on_server_errors(self):
"""
Verify that transient server errors (5xx) trigger retries.
This tests the status_forcelist configuration which specifies
which HTTP status codes should trigger a retry.
"""
from hanzo_insights.request import _build_flags_session, RETRY_STATUS_FORCELIST
session = _build_flags_session()
adapter = session.get_adapter("https://test.posthog.com")
retry = adapter.max_retries
# Verify the status codes that trigger retries
self.assertEqual(
set(retry.status_forcelist),
set(RETRY_STATUS_FORCELIST),
"Should retry on transient server errors",
)
# Verify specific codes are included
self.assertIn(500, retry.status_forcelist)
self.assertIn(502, retry.status_forcelist)
self.assertIn(503, retry.status_forcelist)
self.assertIn(504, retry.status_forcelist)
# Verify rate limits and quota errors are NOT retried
self.assertNotIn(429, retry.status_forcelist)
self.assertNotIn(402, retry.status_forcelist)
def test_flags_session_has_backoff(self):
"""
Verify that retries use exponential backoff to avoid thundering herd.
"""
from hanzo_insights.request import _build_flags_session
session = _build_flags_session()
adapter = session.get_adapter("https://test.posthog.com")
retry = adapter.max_retries
self.assertEqual(
retry.backoff_factor,
0.5,
"Should use 0.5s backoff factor (0.5s, 1s delays)",
)
class TestFlagsSessionRetryIntegration(unittest.TestCase):
"""Integration tests that verify actual retry behavior with a local server."""
def test_retries_on_503_then_succeeds(self):
"""
Verify that 503 errors trigger retries and eventually succeed.
Uses a local HTTP server that fails twice with 503, then succeeds.
This tests the full retry flow including backoff timing.
"""
import threading
from http.server import HTTPServer, BaseHTTPRequestHandler
from socketserver import ThreadingMixIn
from urllib3.util.retry import Retry
from hanzo_insights.request import HTTPAdapterWithSocketOptions, RETRY_STATUS_FORCELIST
request_count = 0
class RetryTestHandler(BaseHTTPRequestHandler):
protocol_version = "HTTP/1.1"
def do_POST(self):
nonlocal request_count
request_count += 1
# Read and discard request body to prevent connection issues
content_length = int(self.headers.get("Content-Length", 0))
if content_length > 0:
self.rfile.read(content_length)
if request_count <= 2:
self.send_response(503)
self.send_header("Content-Type", "application/json")
body = b'{"error": "Service unavailable"}'
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
else:
self.send_response(200)
self.send_header("Content-Type", "application/json")
body = (
b'{"featureFlags": {"test": true}, "featureFlagPayloads": {}}'
)
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
def log_message(self, format, *args):
pass # Suppress logging
# Use ThreadingMixIn for cleaner shutdown
class ThreadedHTTPServer(ThreadingMixIn, HTTPServer):
daemon_threads = True
# Start server on a random available port
server = ThreadedHTTPServer(("127.0.0.1", 0), RetryTestHandler)
port = server.server_address[1]
server_thread = threading.Thread(target=server.serve_forever)
server_thread.daemon = True
server_thread.start()
try:
# Build session with same retry config as _build_flags_session
# but mounted on http:// for local testing
adapter = HTTPAdapterWithSocketOptions(
max_retries=Retry(
total=2,
connect=2,
read=2,
backoff_factor=0.01, # Fast backoff for testing
status_forcelist=RETRY_STATUS_FORCELIST,
allowed_methods=["POST"],
),
)
session = requests.Session()
session.mount("http://", adapter)
response = session.post(
f"http://127.0.0.1:{port}/flags/?v=2",
json={"distinct_id": "user123"},
timeout=5,
)
# Should succeed on 3rd attempt
self.assertEqual(response.status_code, 200)
self.assertEqual(request_count, 3) # 1 initial + 2 retries
finally:
server.shutdown()
server.server_close()
def test_connection_errors_are_retried(self):
"""
Verify that connection errors (no server) trigger retries.
Binds a socket to get a guaranteed available port, then closes it
so connection attempts fail with ConnectionError.
"""
import socket
import time
from urllib3.util.retry import Retry
from hanzo_insights.request import HTTPAdapterWithSocketOptions, RETRY_STATUS_FORCELIST
# Get an available port by binding then closing a socket
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.bind(("127.0.0.1", 0))
port = sock.getsockname()[1]
sock.close() # Port is now available but nothing is listening
adapter = HTTPAdapterWithSocketOptions(
max_retries=Retry(
total=2,
connect=2,
read=2,
backoff_factor=0.05, # Very fast for testing
status_forcelist=RETRY_STATUS_FORCELIST,
allowed_methods=["POST"],
),
)
session = requests.Session()
session.mount("http://", adapter)
start = time.time()
with self.assertRaises(requests.exceptions.ConnectionError):
session.post(
f"http://127.0.0.1:{port}/flags/?v=2",
json={"distinct_id": "user123"},
timeout=1,
)
elapsed = time.time() - start
# With 3 attempts and backoff, should take more than instant
# but less than timeout (confirms retries happened)
self.assertGreater(elapsed, 0.05, "Should have some delay from retries")
@@ -2,7 +2,7 @@ import unittest
from parameterized import parameterized
from posthog import utils
from hanzo_insights import utils
class TestSizeLimitedDict(unittest.TestCase):
@@ -2,7 +2,7 @@ import unittest
from parameterized import parameterized
from posthog.types import (
from hanzo_insights.types import (
FeatureFlag,
FlagMetadata,
FlagReason,
@@ -1,3 +1,4 @@
import sys
import time
import unittest
from dataclasses import dataclass
@@ -12,8 +13,8 @@ from parameterized import parameterized
from pydantic import BaseModel
from pydantic.v1 import BaseModel as BaseModelV1
from posthog import utils
from posthog.types import FeatureFlagResult
from hanzo_insights import utils
from hanzo_insights.types import FeatureFlagResult
TEST_API_KEY = "kOOlRy2QlMY9jHZQv0bKz0FZyazBUoY8Arj0lFVNjs4"
FAKE_TEST_API_KEY = "random_key"
@@ -95,8 +96,8 @@ class TestUtils(unittest.TestCase):
@parameterized.expand(
[
("http://posthog.io/", "http://posthog.io"),
("http://posthog.io", "http://posthog.io"),
("http://hanzo_insights.io/", "http://hanzo_insights.io"),
("http://hanzo_insights.io", "http://hanzo_insights.io"),
("https://example.com/path/", "https://example.com/path"),
("https://example.com/path", "https://example.com/path"),
]
@@ -122,7 +123,9 @@ class TestUtils(unittest.TestCase):
"bar": 2,
"baz": None,
}
assert utils.clean(ModelV1(foo=1, bar="2")) == {"foo": 1, "bar": "2"}
# Pydantic V1 is not compatible with Python 3.14+
if sys.version_info < (3, 14):
assert utils.clean(ModelV1(foo=1, bar="2")) == {"foo": 1, "bar": "2"}
assert utils.clean(NestedModel(foo=ModelV2(foo="1", bar=2, baz="3"))) == {
"foo": {"foo": "1", "bar": 2, "baz": "3"}
}
@@ -123,6 +123,7 @@ class FlagsResponse(TypedDict, total=False):
errorsWhileComputingFlags: bool
requestId: str
quotaLimit: Optional[List[str]]
evaluatedAt: Optional[int]
class FlagsAndPayloads(TypedDict, total=True):
@@ -306,3 +307,42 @@ def to_payloads(response: FlagsResponse) -> Optional[dict[str, str]]:
and value.enabled
and value.metadata.payload is not None
}
class FeatureFlagError:
"""Error type constants for the $feature_flag_error property.
These values are sent in analytics events to track flag evaluation failures.
They should not be changed without considering impact on existing dashboards
and queries that filter on these values.
Error values:
ERRORS_WHILE_COMPUTING: Server returned errorsWhileComputingFlags=true
FLAG_MISSING: Requested flag not in API response
QUOTA_LIMITED: Rate/quota limit exceeded
TIMEOUT: Request timed out
CONNECTION_ERROR: Network connectivity issue
UNKNOWN_ERROR: Unexpected exceptions
For API errors with status codes, use the api_error() method which returns
a string like "api_error_500".
"""
ERRORS_WHILE_COMPUTING = "errors_while_computing_flags"
FLAG_MISSING = "flag_missing"
QUOTA_LIMITED = "quota_limited"
TIMEOUT = "timeout"
CONNECTION_ERROR = "connection_error"
UNKNOWN_ERROR = "unknown_error"
@staticmethod
def api_error(status: Union[int, str]) -> str:
"""Generate API error string with status code.
Args:
status: HTTP status code from the API error
Returns:
Error string like "api_error_500"
"""
return f"api_error_{status}"
+2 -2
View File
@@ -16,7 +16,7 @@ import distro # For Linux OS detection
import six
from dateutil.tz import tzlocal, tzutc
log = logging.getLogger("posthog")
log = logging.getLogger("hanzo_insights")
def is_naive(dt):
@@ -277,7 +277,7 @@ class FlagCache:
class RedisFlagCache:
def __init__(
self, redis_client, default_ttl=300, stale_ttl=3600, key_prefix="posthog:flags:"
self, redis_client, default_ttl=300, stale_ttl=3600, key_prefix="insights:flags:"
):
self.redis = redis_client
self.default_ttl = default_ttl
+1
View File
@@ -0,0 +1 @@
VERSION = "7.9.7"
+3
View File
@@ -0,0 +1,3 @@
# Convenience re-export so `from insights import Insights` works.
from hanzo_insights import * # noqa: F401, F403
from hanzo_insights import Insights, Client # noqa: F401
@@ -1,10 +1,10 @@
"""
Test that verifies exception capture functionality.
These tests verify that exceptions are actually captured to PostHog, not just that
These tests verify that exceptions are actually captured to Insights, not just that
500 responses are returned.
Without process_exception(), view exceptions are NOT captured to PostHog (v6.7.11 and earlier).
Without process_exception(), view exceptions are NOT captured to Insights (v6.7.11 and earlier).
With process_exception(), Django calls this method to capture exceptions before
converting them to 500 responses.
"""
@@ -30,7 +30,7 @@ def asgi_app():
@pytest.mark.asyncio
async def test_async_exception_is_captured(asgi_app):
"""
Test that async view exceptions are captured to PostHog.
Test that async view exceptions are captured to Insights.
The middleware's process_exception() method ensures exceptions are captured.
Without it (v6.7.11 and earlier), exceptions are NOT captured even though 500 is returned.
@@ -50,8 +50,8 @@ async def test_async_exception_is_captured(asgi_app):
}
)
# Patch at the posthog module level where middleware imports from
with patch("posthog.capture_exception", side_effect=mock_capture):
# Patch at the hanzo_insights module level where middleware imports from
with patch("hanzo_insights.capture_exception", side_effect=mock_capture):
async with AsyncClient(
transport=ASGITransport(app=asgi_app), base_url="http://testserver"
) as ac:
@@ -60,8 +60,8 @@ async def test_async_exception_is_captured(asgi_app):
# Django returns 500
assert response.status_code == 500
# CRITICAL: Verify PostHog captured the exception
assert len(captured) > 0, "Exception was NOT captured to PostHog!"
# CRITICAL: Verify Insights captured the exception
assert len(captured) > 0, "Exception was NOT captured to Insights!"
# Verify it's the right exception
exception_data = captured[0]
@@ -72,7 +72,7 @@ async def test_async_exception_is_captured(asgi_app):
@pytest.mark.asyncio
async def test_sync_exception_is_captured(asgi_app):
"""
Test that sync view exceptions are captured to PostHog.
Test that sync view exceptions are captured to Insights.
The middleware's process_exception() method ensures exceptions are captured.
Without it (v6.7.11 and earlier), exceptions are NOT captured even though 500 is returned.
@@ -92,8 +92,8 @@ async def test_sync_exception_is_captured(asgi_app):
}
)
# Patch at the posthog module level where middleware imports from
with patch("posthog.capture_exception", side_effect=mock_capture):
# Patch at the hanzo_insights module level where middleware imports from
with patch("hanzo_insights.capture_exception", side_effect=mock_capture):
async with AsyncClient(
transport=ASGITransport(app=asgi_app), base_url="http://testserver"
) as ac:
@@ -102,8 +102,8 @@ async def test_sync_exception_is_captured(asgi_app):
# Django returns 500
assert response.status_code == 500
# CRITICAL: Verify PostHog captured the exception
assert len(captured) > 0, "Exception was NOT captured to PostHog!"
# CRITICAL: Verify Insights captured the exception
assert len(captured) > 0, "Exception was NOT captured to Insights!"
# Verify it's the right exception
exception_data = captured[0]
+5 -5
View File
@@ -1,5 +1,5 @@
"""
Tests for PostHog Django middleware in async context.
Tests for Insights Django middleware in async context.
These tests verify that the middleware correctly handles:
1. Async user access (request.auser() in Django 5)
@@ -103,7 +103,7 @@ async def test_async_authenticated_user_access(asgi_app):
# Make request with session cookie - this should trigger the bug in v6.7.11
# Disable exception capture to see the SynchronousOnlyOperation clearly
with override_settings(POSTHOG_MW_CAPTURE_EXCEPTIONS=False):
with override_settings(INSIGHTS_MW_CAPTURE_EXCEPTIONS=False):
async with AsyncClient(
transport=ASGITransport(app=asgi_app),
base_url="http://testserver",
@@ -139,10 +139,10 @@ async def test_async_exception_capture(asgi_app):
"""
Test that middleware handles exceptions from async views.
The middleware's process_exception() method captures view exceptions to PostHog
The middleware's process_exception() method captures view exceptions to Insights
before Django converts them to 500 responses. This test verifies the exception
causes a 500 response. See test_exception_capture.py for tests that verify
actual exception capture to PostHog.
actual exception capture to Insights.
"""
async with AsyncClient(
transport=ASGITransport(app=asgi_app), base_url="http://testserver"
@@ -158,7 +158,7 @@ async def test_sync_exception_capture(asgi_app):
"""
Test that middleware handles exceptions from sync views.
The middleware's process_exception() method captures view exceptions to PostHog.
The middleware's process_exception() method captures view exceptions to Insights.
This test verifies the exception causes a 500 response.
"""
async with AsyncClient(
@@ -47,7 +47,7 @@ MIDDLEWARE = [
"django.contrib.auth.middleware.AuthenticationMiddleware",
"django.contrib.messages.middleware.MessageMiddleware",
"django.middleware.clickjacking.XFrameOptionsMiddleware",
"posthog.integrations.django.PosthogContextMiddleware", # Test PostHog middleware
"hanzo_insights.integrations.django.InsightsContextMiddleware", # Test Insights middleware
]
ROOT_URLCONF = "testdjango.urls"
@@ -123,7 +123,7 @@ STATIC_URL = "static/"
DEFAULT_AUTO_FIELD = "django.db.models.BigAutoField"
# PostHog settings for testing
POSTHOG_API_KEY = "test-key"
# Insights settings for testing
INSIGHTS_API_KEY = "test-key"
POSTHOG_HOST = "https://app.posthog.com"
POSTHOG_MW_CAPTURE_EXCEPTIONS = True
INSIGHTS_MW_CAPTURE_EXCEPTIONS = True
@@ -1,5 +1,5 @@
"""
Test views for validating PostHog middleware with Django 5 ASGI.
Test views for validating Insights middleware with Django 5 ASGI.
"""
from django.http import JsonResponse
-5
View File
@@ -26,14 +26,9 @@ posthog/client.py:0: error: Incompatible types in assignment (expression has typ
posthog/client.py:0: error: Incompatible types in assignment (expression has type "dict[Any, Any]", variable has type "None") [assignment]
posthog/client.py:0: error: "None" has no attribute "__iter__" (not iterable) [attr-defined]
posthog/client.py:0: error: Statement is unreachable [unreachable]
posthog/client.py:0: error: Incompatible types in assignment (expression has type "Any | dict[Any, Any]", variable has type "None") [assignment]
posthog/client.py:0: error: Incompatible types in assignment (expression has type "Any | dict[Any, Any]", variable has type "None") [assignment]
posthog/client.py:0: error: Incompatible types in assignment (expression has type "dict[Never, Never]", variable has type "None") [assignment]
posthog/client.py:0: error: Incompatible types in assignment (expression has type "dict[Never, Never]", variable has type "None") [assignment]
posthog/client.py:0: error: Right operand of "and" is never evaluated [unreachable]
posthog/client.py:0: error: Incompatible types in assignment (expression has type "Poller", variable has type "None") [assignment]
posthog/client.py:0: error: "None" has no attribute "start" [attr-defined]
posthog/client.py:0: error: "None" has no attribute "get" [attr-defined]
posthog/client.py:0: error: Statement is unreachable [unreachable]
posthog/client.py:0: error: Statement is unreachable [unreachable]
posthog/client.py:0: error: Name "urlparse" already defined (possibly by an import) [no-redef]
+2 -2
View File
@@ -17,10 +17,10 @@ ignore_missing_imports = True
[mypy-sentry_sdk.*]
ignore_missing_imports = True
[mypy-posthog.test.*]
[mypy-hanzo_insights.test.*]
ignore_errors = True
[mypy-posthog.*.test.*]
[mypy-hanzo_insights.*.test.*]
ignore_errors = True
[mypy-openai.*]
-589
View File
@@ -1,589 +0,0 @@
import time
import uuid
from typing import Any, Callable, Dict, List, Optional, cast
from posthog.client import Client as PostHogClient
from posthog.ai.types import FormattedMessage, StreamingEventData, TokenUsage
from posthog.ai.sanitization import (
sanitize_openai,
sanitize_anthropic,
sanitize_gemini,
sanitize_langchain,
)
def merge_usage_stats(
target: TokenUsage, source: TokenUsage, mode: str = "incremental"
) -> None:
"""
Merge streaming usage statistics into target dict, handling None values.
Supports two modes:
- "incremental": Add source values to target (for APIs that report new tokens)
- "cumulative": Replace target with source values (for APIs that report totals)
Args:
target: Dictionary to update with usage stats
source: TokenUsage that may contain None values
mode: Either "incremental" or "cumulative"
"""
if mode == "incremental":
# Add new values to existing totals
source_input = source.get("input_tokens")
if source_input is not None:
current = target.get("input_tokens") or 0
target["input_tokens"] = current + source_input
source_output = source.get("output_tokens")
if source_output is not None:
current = target.get("output_tokens") or 0
target["output_tokens"] = current + source_output
source_cache_read = source.get("cache_read_input_tokens")
if source_cache_read is not None:
current = target.get("cache_read_input_tokens") or 0
target["cache_read_input_tokens"] = current + source_cache_read
source_cache_creation = source.get("cache_creation_input_tokens")
if source_cache_creation is not None:
current = target.get("cache_creation_input_tokens") or 0
target["cache_creation_input_tokens"] = current + source_cache_creation
source_reasoning = source.get("reasoning_tokens")
if source_reasoning is not None:
current = target.get("reasoning_tokens") or 0
target["reasoning_tokens"] = current + source_reasoning
source_web_search = source.get("web_search_count")
if source_web_search is not None:
current = target.get("web_search_count") or 0
target["web_search_count"] = max(current, source_web_search)
elif mode == "cumulative":
# Replace with latest values (already cumulative)
if source.get("input_tokens") is not None:
target["input_tokens"] = source["input_tokens"]
if source.get("output_tokens") is not None:
target["output_tokens"] = source["output_tokens"]
if source.get("cache_read_input_tokens") is not None:
target["cache_read_input_tokens"] = source["cache_read_input_tokens"]
if source.get("cache_creation_input_tokens") is not None:
target["cache_creation_input_tokens"] = source[
"cache_creation_input_tokens"
]
if source.get("reasoning_tokens") is not None:
target["reasoning_tokens"] = source["reasoning_tokens"]
if source.get("web_search_count") is not None:
target["web_search_count"] = source["web_search_count"]
else:
raise ValueError(f"Invalid mode: {mode}. Must be 'incremental' or 'cumulative'")
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) -> TokenUsage:
"""
Extract usage statistics from response based on provider.
Delegates to provider-specific converter functions.
"""
if provider == "anthropic":
from posthog.ai.anthropic.anthropic_converter import (
extract_anthropic_usage_from_response,
)
return extract_anthropic_usage_from_response(response)
elif provider == "openai":
from posthog.ai.openai.openai_converter import (
extract_openai_usage_from_response,
)
return extract_openai_usage_from_response(response)
elif provider == "gemini":
from posthog.ai.gemini.gemini_converter import (
extract_gemini_usage_from_response,
)
return extract_gemini_usage_from_response(response)
return TokenUsage(input_tokens=0, output_tokens=0)
def format_response(response, provider: str):
"""
Format a regular (non-streaming) response.
"""
if provider == "anthropic":
from posthog.ai.anthropic.anthropic_converter import format_anthropic_response
return format_anthropic_response(response)
elif provider == "openai":
from posthog.ai.openai.openai_converter import format_openai_response
return format_openai_response(response)
elif provider == "gemini":
from posthog.ai.gemini.gemini_converter import format_gemini_response
return format_gemini_response(response)
return []
def extract_available_tool_calls(provider: str, kwargs: Dict[str, Any]):
"""
Extract available tool calls for the given provider.
"""
if provider == "anthropic":
from posthog.ai.anthropic.anthropic_converter import extract_anthropic_tools
return extract_anthropic_tools(kwargs)
elif provider == "gemini":
from posthog.ai.gemini.gemini_converter import extract_gemini_tools
return extract_gemini_tools(kwargs)
elif provider == "openai":
from posthog.ai.openai.openai_converter import extract_openai_tools
return extract_openai_tools(kwargs)
return None
def merge_system_prompt(
kwargs: Dict[str, Any], provider: str
) -> List[FormattedMessage]:
"""
Merge system prompts and format messages for the given provider.
"""
if provider == "anthropic":
from posthog.ai.anthropic.anthropic_converter import format_anthropic_input
messages = kwargs.get("messages") or []
system = kwargs.get("system")
return format_anthropic_input(messages, system)
elif provider == "gemini":
from posthog.ai.gemini.gemini_converter import format_gemini_input_with_system
contents = kwargs.get("contents", [])
config = kwargs.get("config")
return format_gemini_input_with_system(contents, config)
elif provider == "openai":
from posthog.ai.openai.openai_converter import format_openai_input
# For OpenAI, handle both Chat Completions and Responses API
messages_param = kwargs.get("messages")
input_param = kwargs.get("input")
# Get base formatted messages
messages = format_openai_input(messages_param, input_param)
# Check if system prompt is provided as a separate parameter
if kwargs.get("system") is not None:
has_system = any(msg.get("role") == "system" for msg in messages)
if not has_system:
system_msg = cast(
FormattedMessage,
{"role": "system", "content": kwargs.get("system")},
)
messages = [system_msg] + messages
# For Responses API, add instructions to the system prompt if provided
if kwargs.get("instructions") is not None:
# Find the system message if it exists
system_idx = next(
(i for i, msg in enumerate(messages) if msg.get("role") == "system"),
None,
)
if system_idx is not None:
# Append instructions to existing system message
system_content = messages[system_idx].get("content", "")
messages[system_idx]["content"] = (
f"{system_content}\n\n{kwargs.get('instructions')}"
)
else:
# Create a new system message with instructions
instruction_msg = cast(
FormattedMessage,
{"role": "system", "content": kwargs.get("instructions")},
)
messages = [instruction_msg] + messages
return messages
# Default case - return empty list
return []
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: str,
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: TokenUsage = TokenUsage()
error_params: Dict[str, Any] = {}
try:
response = call_method(**kwargs)
except Exception as exc:
error = exc
http_status = getattr(
exc, "status_code", 0
) # default to 0 becuase its likely an SDK error
error_params = {
"$ai_is_error": True,
"$ai_error": exc.__str__(),
}
finally:
end_time = time.time()
latency = end_time - start_time
if posthog_trace_id is None:
posthog_trace_id = str(uuid.uuid4())
if response and (
hasattr(response, "usage")
or (provider == "gemini" and hasattr(response, "usage_metadata"))
):
usage = get_usage(response, provider)
messages = merge_system_prompt(kwargs, provider)
sanitized_messages = sanitize_messages(messages, 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, sanitized_messages
),
"$ai_output_choices": with_privacy_mode(
ph_client, posthog_privacy_mode, format_response(response, provider)
),
"$ai_http_status": http_status,
"$ai_input_tokens": usage.get("input_tokens", 0),
"$ai_output_tokens": usage.get("output_tokens", 0),
"$ai_latency": latency,
"$ai_trace_id": posthog_trace_id,
"$ai_base_url": str(base_url),
**(posthog_properties or {}),
**(error_params or {}),
}
available_tool_calls = extract_available_tool_calls(provider, kwargs)
if available_tool_calls:
event_properties["$ai_tools"] = available_tool_calls
cache_read = usage.get("cache_read_input_tokens")
if cache_read is not None and cache_read > 0:
event_properties["$ai_cache_read_input_tokens"] = cache_read
cache_creation = usage.get("cache_creation_input_tokens")
if cache_creation is not None and cache_creation > 0:
event_properties["$ai_cache_creation_input_tokens"] = cache_creation
reasoning = usage.get("reasoning_tokens")
if reasoning is not None and reasoning > 0:
event_properties["$ai_reasoning_tokens"] = reasoning
web_search_count = usage.get("web_search_count")
if web_search_count is not None and web_search_count > 0:
event_properties["$ai_web_search_count"] = web_search_count
if posthog_distinct_id is None:
event_properties["$process_person_profile"] = False
# Process instructions for Responses API
if provider == "openai" and kwargs.get("instructions") is not None:
event_properties["$ai_instructions"] = with_privacy_mode(
ph_client, posthog_privacy_mode, kwargs.get("instructions")
)
# send the event to posthog
if hasattr(ph_client, "capture") and callable(ph_client.capture):
ph_client.capture(
distinct_id=posthog_distinct_id or posthog_trace_id,
event="$ai_generation",
properties=event_properties,
groups=posthog_groups,
)
if error:
raise error
return response
async def call_llm_and_track_usage_async(
posthog_distinct_id: Optional[str],
ph_client: PostHogClient,
provider: str,
posthog_trace_id: Optional[str],
posthog_properties: Optional[Dict[str, Any]],
posthog_privacy_mode: bool,
posthog_groups: Optional[Dict[str, Any]],
base_url: str,
call_async_method: Callable[..., Any],
**kwargs: Any,
) -> Any:
start_time = time.time()
response = None
error = None
http_status = 200
usage: TokenUsage = TokenUsage()
error_params: Dict[str, Any] = {}
try:
response = await call_async_method(**kwargs)
except Exception as exc:
error = exc
http_status = getattr(
exc, "status_code", 0
) # default to 0 because its likely an SDK error
error_params = {
"$ai_is_error": True,
"$ai_error": exc.__str__(),
}
finally:
end_time = time.time()
latency = end_time - start_time
if posthog_trace_id is None:
posthog_trace_id = str(uuid.uuid4())
if response and (
hasattr(response, "usage")
or (provider == "gemini" and hasattr(response, "usage_metadata"))
):
usage = get_usage(response, provider)
messages = merge_system_prompt(kwargs, provider)
sanitized_messages = sanitize_messages(messages, 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, sanitized_messages
),
"$ai_output_choices": with_privacy_mode(
ph_client, posthog_privacy_mode, format_response(response, provider)
),
"$ai_http_status": http_status,
"$ai_input_tokens": usage.get("input_tokens", 0),
"$ai_output_tokens": usage.get("output_tokens", 0),
"$ai_latency": latency,
"$ai_trace_id": posthog_trace_id,
"$ai_base_url": str(base_url),
**(posthog_properties or {}),
**(error_params or {}),
}
available_tool_calls = extract_available_tool_calls(provider, kwargs)
if available_tool_calls:
event_properties["$ai_tools"] = available_tool_calls
cache_read = usage.get("cache_read_input_tokens")
if cache_read is not None and cache_read > 0:
event_properties["$ai_cache_read_input_tokens"] = cache_read
cache_creation = usage.get("cache_creation_input_tokens")
if cache_creation is not None and cache_creation > 0:
event_properties["$ai_cache_creation_input_tokens"] = cache_creation
reasoning = usage.get("reasoning_tokens")
if reasoning is not None and reasoning > 0:
event_properties["$ai_reasoning_tokens"] = reasoning
web_search_count = usage.get("web_search_count")
if web_search_count is not None and web_search_count > 0:
event_properties["$ai_web_search_count"] = web_search_count
if posthog_distinct_id is None:
event_properties["$process_person_profile"] = False
# Process instructions for Responses API
if provider == "openai" and kwargs.get("instructions") is not None:
event_properties["$ai_instructions"] = with_privacy_mode(
ph_client, posthog_privacy_mode, kwargs.get("instructions")
)
# send the event to posthog
if hasattr(ph_client, "capture") and callable(ph_client.capture):
ph_client.capture(
distinct_id=posthog_distinct_id or posthog_trace_id,
event="$ai_generation",
properties=event_properties,
groups=posthog_groups,
)
if error:
raise error
return response
def sanitize_messages(data: Any, provider: str) -> Any:
"""Sanitize messages using provider-specific sanitization functions."""
if provider == "anthropic":
return sanitize_anthropic(data)
elif provider == "openai":
return sanitize_openai(data)
elif provider == "gemini":
return sanitize_gemini(data)
elif provider == "langchain":
return sanitize_langchain(data)
return data
def with_privacy_mode(ph_client: PostHogClient, privacy_mode: bool, value: Any):
if ph_client.privacy_mode or privacy_mode:
return None
return value
def capture_streaming_event(
ph_client: PostHogClient,
event_data: StreamingEventData,
):
"""
Unified streaming event capture for all LLM providers.
This function handles the common logic for capturing streaming events across all providers.
All provider-specific formatting should be done BEFORE calling this function.
The function handles:
- Building PostHog event properties
- Extracting and adding tools based on provider
- Applying privacy mode
- Adding special token fields (cache, reasoning)
- Provider-specific fields (e.g., OpenAI instructions)
- Sending the event to PostHog
Args:
ph_client: PostHog client instance
event_data: Standardized streaming event data containing all necessary information
"""
trace_id = event_data.get("trace_id") or str(uuid.uuid4())
# Build base event properties
event_properties = {
"$ai_provider": event_data["provider"],
"$ai_model": event_data["model"],
"$ai_model_parameters": get_model_params(event_data["kwargs"]),
"$ai_input": with_privacy_mode(
ph_client,
event_data["privacy_mode"],
event_data["formatted_input"],
),
"$ai_output_choices": with_privacy_mode(
ph_client,
event_data["privacy_mode"],
event_data["formatted_output"],
),
"$ai_http_status": 200,
"$ai_input_tokens": event_data["usage_stats"].get("input_tokens", 0),
"$ai_output_tokens": event_data["usage_stats"].get("output_tokens", 0),
"$ai_latency": event_data["latency"],
"$ai_trace_id": trace_id,
"$ai_base_url": str(event_data["base_url"]),
**(event_data.get("properties") or {}),
}
# Extract and add tools based on provider
available_tools = extract_available_tool_calls(
event_data["provider"],
event_data["kwargs"],
)
if available_tools:
event_properties["$ai_tools"] = available_tools
# Add optional token fields
# For Anthropic, always include cache fields even if 0 (backward compatibility)
# For others, only include if present and non-zero
if event_data["provider"] == "anthropic":
# Anthropic always includes cache fields
cache_read = event_data["usage_stats"].get("cache_read_input_tokens", 0)
cache_creation = event_data["usage_stats"].get("cache_creation_input_tokens", 0)
event_properties["$ai_cache_read_input_tokens"] = cache_read
event_properties["$ai_cache_creation_input_tokens"] = cache_creation
else:
# Other providers only include if non-zero
optional_token_fields = [
"cache_read_input_tokens",
"cache_creation_input_tokens",
"reasoning_tokens",
]
for field in optional_token_fields:
value = event_data["usage_stats"].get(field)
if value is not None and isinstance(value, int) and value > 0:
event_properties[f"$ai_{field}"] = value
# Add web search count if present (all providers)
web_search_count = event_data["usage_stats"].get("web_search_count")
if (
web_search_count is not None
and isinstance(web_search_count, int)
and web_search_count > 0
):
event_properties["$ai_web_search_count"] = web_search_count
# Handle provider-specific fields
if (
event_data["provider"] == "openai"
and event_data["kwargs"].get("instructions") is not None
):
event_properties["$ai_instructions"] = with_privacy_mode(
ph_client,
event_data["privacy_mode"],
event_data["kwargs"]["instructions"],
)
if event_data.get("distinct_id") is None:
event_properties["$process_person_profile"] = False
# Send event to PostHog
if hasattr(ph_client, "capture"):
ph_client.capture(
distinct_id=event_data.get("distinct_id") or trace_id,
event="$ai_generation",
properties=event_properties,
groups=event_data.get("groups"),
)

Some files were not shown because too many files have changed in this diff Show More