From a473b998dcde751d957dd95f79b0d15381af9055 Mon Sep 17 00:00:00 2001 From: Eli Reisman Date: Tue, 6 Oct 2026 20:23:38 -0700 Subject: [PATCH] feat!: AI wrappers always send through capture_ai enable_full_ai_capture now controls only truncation and media redaction. A client object without capture_ai still gets capture calls. --- posthog/__init__.py | 4 +- posthog/ai/utils.py | 15 +- posthog/client.py | 10 +- posthog/test/ai/anthropic/test_anthropic.py | 149 +++++----- .../ai/claude_agent_sdk/test_processor.py | 50 ++-- posthog/test/ai/gemini/test_gemini.py | 124 ++++---- posthog/test/ai/gemini/test_gemini_async.py | 80 +++--- posthog/test/ai/gemini/test_gemini_parity.py | 4 +- posthog/test/ai/langchain/test_callbacks.py | 270 +++++++++--------- posthog/test/ai/langchain/test_middleware.py | 6 +- posthog/test/ai/openai/test_async_parity.py | 22 +- posthog/test/ai/openai/test_openai.py | 200 ++++++------- .../test/ai/openai_agents/test_processor.py | 105 +++---- posthog/test/ai/test_token_reporting.py | 4 +- posthog/test/test_ai_capture_lane.py | 83 +----- 15 files changed, 486 insertions(+), 640 deletions(-) diff --git a/posthog/__init__.py b/posthog/__init__.py index d228d811..027d6425 100644 --- a/posthog/__init__.py +++ b/posthog/__init__.py @@ -485,8 +485,8 @@ def get_context_options() -> Dict[str, Any]: # We recommend setting this to False if you are only using the personalApiKey for evaluating remote config payloads via `get_remote_config_payload` and not using local evaluation. enable_local_evaluation = True # type: bool flag_definition_cache_provider = None # type: Optional[FlagDefinitionCacheProvider] -# Routes AI SDK wrapper events through the dedicated AI capture lane, skips -# truncation, and passes media unredacted. `privacy_mode` always wins. +# Captures full AI content in AI SDK wrapper events: skips truncation and +# passes media unredacted. `privacy_mode` always wins. enable_full_ai_capture = False # type: bool # Deprecated aliases for `enable_full_ai_capture`. _use_ai_lane = False # type: bool diff --git a/posthog/ai/utils.py b/posthog/ai/utils.py index ae652e80..2a6d20cc 100644 --- a/posthog/ai/utils.py +++ b/posthog/ai/utils.py @@ -5,7 +5,7 @@ from posthog import get_tags, identify_context, new_context, tag, contexts from posthog.ai.gateway import warn_if_posthog_ai_gateway -from posthog.ai.sanitization import _full_ai_capture_enabled, redact_media +from posthog.ai.sanitization import redact_media from posthog.ai.sanitization import sanitize_messages # noqa: F401 -- re-exported for back-compat from posthog.ai.types import FormattedMessage, StreamingEventData, TokenUsage from posthog.client import Client as PostHogClient @@ -64,11 +64,6 @@ def _get_tokens_source( return "sdk" -def _ai_lane_enabled(ph_client) -> bool: - """The client's full-AI-capture opt-in routes wrapper events onto the AI lane.""" - return _full_ai_capture_enabled(ph_client) - - def _capture_ai_event(ph_client, event: str, *, personless: bool = False, **kwargs): """Capture a wrapper-emitted AI event with the PostHog AI library identity. @@ -87,10 +82,10 @@ def _capture_ai_event(ph_client, event: str, *, personless: bool = False, **kwar **(kwargs.get("options") or {}), "process_person_profile": False, } - if _ai_lane_enabled(ph_client): - capture_ai = getattr(ph_client, "capture_ai", None) - if callable(capture_ai): - return capture_ai(event=event, **kwargs) + # Clients without capture_ai, such as custom stand-ins, keep the analytics lane. + capture_ai = getattr(ph_client, "capture_ai", None) + if callable(capture_ai): + return capture_ai(event=event, **kwargs) return ph_client.capture(event=event, **kwargs) diff --git a/posthog/client.py b/posthog/client.py index 1ce5108b..0f7b415e 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -846,11 +846,11 @@ def __init__( captured exceptions. Defaults to the current working directory. privacy_mode: For AI observability, capture usage metadata without prompt inputs or outputs. - enable_full_ai_capture: Route PostHog AI wrapper events through - the dedicated AI capture endpoint and capture full AI content: - skips string truncation and passes media (base64/data URIs) - through unredacted. ``privacy_mode`` always wins. Defaults to - False. + enable_full_ai_capture: Capture full AI content in PostHog AI + wrapper events: skips string truncation and passes media + (base64/data URIs) through unredacted. ``privacy_mode`` always + wins. Wrapper events use the AI capture endpoint whether or not + this is set. Defaults to False. before_send: Optional callback that can modify or drop events before upload. Return ``None`` to drop an event. Context tags, context options, ``super_properties``, ``super_options`` and diff --git a/posthog/test/ai/anthropic/test_anthropic.py b/posthog/test/ai/anthropic/test_anthropic.py index a4a3e9bc..2774933a 100644 --- a/posthog/test/ai/anthropic/test_anthropic.py +++ b/posthog/test/ai/anthropic/test_anthropic.py @@ -430,9 +430,9 @@ def test_basic_completion(mock_client, mock_anthropic_response): ) assert response == mock_anthropic_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -476,7 +476,7 @@ def test_tokens_source_passthrough(mock_client, mock_anthropic_response): posthog_properties={"$ai_input_tokens": 99999}, ) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_tokens_source"] == "passthrough" assert props["$ai_input_tokens"] == 99999 @@ -494,9 +494,9 @@ def test_groups(mock_client, mock_anthropic_response): ) assert response == mock_anthropic_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] assert call_args["groups"] == {"company": "test_company"} @@ -513,9 +513,9 @@ def test_privacy_mode_local(mock_client, mock_anthropic_response): ) assert response == mock_anthropic_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] is None assert props["$ai_output_choices"] is None @@ -535,9 +535,9 @@ def test_privacy_mode_global(mock_client, mock_anthropic_response): ) assert response == mock_anthropic_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] is None assert props["$ai_output_choices"] is None @@ -561,9 +561,9 @@ def test_basic_integration(mock_client): system="You must always answer with 'Bar'.", ) - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" assert call_args["event"] == "$ai_generation" @@ -607,9 +607,9 @@ async def mock_async_create(**kwargs): posthog_properties={"foo": "bar"}, ) - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -669,9 +669,9 @@ async def async_create_wrapper(**kwargs): [c async for c in response] # Capture happens in the async finally block before generator completes - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] == [ @@ -691,9 +691,9 @@ def test_error(mock_client, mock_anthropic_response): messages=[{"role": "user", "content": "Hello"}], ) - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_is_error"] is True assert props["$ai_error"] == "Test error" @@ -713,9 +713,9 @@ def test_cached_tokens(mock_client, mock_anthropic_response_with_cached_tokens): ) assert response == mock_anthropic_response_with_cached_tokens - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -754,7 +754,7 @@ def test_preserves_cache_creation_ttl_breakdown_non_streaming( ) assert_cache_creation_ttl_breakdown_preserved( - mock_client.capture.call_args.kwargs["properties"] + mock_client.capture_ai.call_args.kwargs["properties"] ) @@ -776,7 +776,7 @@ async def mock_async_create(**kwargs): ) assert_cache_creation_ttl_breakdown_preserved( - mock_client.capture.call_args.kwargs["properties"] + mock_client.capture_ai.call_args.kwargs["properties"] ) @@ -794,7 +794,7 @@ def test_preserves_cache_creation_ttl_breakdown_streaming(mock_client): list(response) assert_cache_creation_ttl_breakdown_preserved( - mock_client.capture.call_args.kwargs["properties"] + mock_client.capture_ai.call_args.kwargs["properties"] ) @@ -820,7 +820,7 @@ async def mock_async_create(**kwargs): [event async for event in response] assert_cache_creation_ttl_breakdown_preserved( - mock_client.capture.call_args.kwargs["properties"] + mock_client.capture_ai.call_args.kwargs["properties"] ) @@ -843,7 +843,7 @@ def test_anthropic_bedrock_preserves_cache_creation_ttl_breakdown( ) assert_cache_creation_ttl_breakdown_preserved( - mock_client.capture.call_args.kwargs["properties"] + mock_client.capture_ai.call_args.kwargs["properties"] ) @@ -882,9 +882,9 @@ def test_tool_definition(mock_client, mock_anthropic_response): ) assert response == mock_anthropic_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -936,9 +936,9 @@ def test_tool_calls_in_output_choices( ) assert response == mock_anthropic_response_with_tool_calls - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -999,9 +999,9 @@ def test_tool_calls_only_no_content( ) assert response == mock_anthropic_response_tool_calls_only - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -1068,9 +1068,9 @@ async def run_test(): response = asyncio.run(run_test()) assert response == mock_anthropic_response_with_tool_calls - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -1136,9 +1136,9 @@ def test_streaming_with_tool_calls(mock_client, mock_anthropic_stream_with_tools list(response) # Capture happens synchronously when generator is exhausted - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -1278,7 +1278,7 @@ def stream_generator(): ) list(response) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_output_choices"] == [ { "role": "assistant", @@ -1345,9 +1345,9 @@ async def run_test(): asyncio.run(run_test()) # Capture completes before asyncio.run() returns - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -1427,9 +1427,9 @@ def test_streaming_with_thinking(mock_client, mock_anthropic_stream_with_thinkin # Consume the stream - this triggers the finally block synchronously list(response) - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] output_choices = props["$ai_output_choices"] @@ -1479,9 +1479,9 @@ async def run_test(): asyncio.run(run_test()) - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] output_choices = props["$ai_output_choices"] @@ -1533,9 +1533,9 @@ def __init__(self): ) assert response == mock_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify web search count is captured @@ -1615,9 +1615,9 @@ def test_streaming_with_web_search(mock_client, mock_anthropic_stream_with_web_s list(response) # Capture happens synchronously when generator is exhausted - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify web search count is captured @@ -1672,9 +1672,9 @@ async def run_test(): response = asyncio.run(run_test()) assert response == mock_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify web search count is captured @@ -1719,9 +1719,9 @@ async def run_test(): asyncio.run(run_test()) # Capture completes before asyncio.run() returns - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify web search count is captured @@ -1778,7 +1778,7 @@ async def create(**kwargs): if stream: list(response) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] assert call_args["distinct_id"] == "trace-123" assert call_args["options"] == {"process_person_profile": False} @@ -1798,7 +1798,7 @@ def test_explicit_distinct_id_creates_person_profile( posthog_trace_id="trace-123", ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "user-123" @@ -1820,7 +1820,7 @@ def test_outer_context_distinct_id_is_used(mock_client, mock_anthropic_response) posthog_trace_id="trace-123", ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "outer-user-456" @@ -1845,7 +1845,7 @@ def test_explicit_distinct_id_overrides_outer_context( posthog_trace_id="trace-123", ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] assert call_args["distinct_id"] == "explicit-user-789" @@ -1862,7 +1862,7 @@ def test_integration_stop_reason(mock_client): posthog_distinct_id="test-id", ) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_stop_reason"] in ("end_turn", "max_tokens") assert props["$ai_provider"] == "anthropic" assert props["$ai_input_tokens"] > 0 @@ -1963,9 +1963,9 @@ async def test_streaming_sync_async_accumulation_parity(mock_client): 0, ) list(sync_response) - sync_properties = mock_client.capture.call_args.kwargs["properties"] + sync_properties = mock_client.capture_ai.call_args.kwargs["properties"] - mock_client.capture.reset_mock() + mock_client.capture_ai.reset_mock() async_client = AsyncAnthropic(api_key="test-key", posthog_client=mock_client) async_response = async_client.messages._track_streaming_response( @@ -1979,7 +1979,7 @@ async def test_streaming_sync_async_accumulation_parity(mock_client): 0, ) [event async for event in async_response] - async_properties = mock_client.capture.call_args.kwargs["properties"] + async_properties = mock_client.capture_ai.call_args.kwargs["properties"] parity_keys = ( "$ai_input", @@ -2019,9 +2019,9 @@ async def test_streaming_sync_async_exceptions_capture_and_close(mock_client): list(stream) assert sync_source.closed is True - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - mock_client.capture.reset_mock() + mock_client.capture_ai.reset_mock() async_source = FailingRecordingAsyncStream([]) async_client = AsyncAnthropic(api_key="test-key", posthog_client=mock_client) @@ -2040,7 +2040,7 @@ async def test_streaming_sync_async_exceptions_capture_and_close(mock_client): [event async for event in stream] assert async_source.closed is True - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 def test_messages_stream_preserves_native_manager_helpers_close_and_tracking( @@ -2064,8 +2064,8 @@ def test_messages_stream_preserves_native_manager_helpers_close_and_tracking( assert text == ["Hi"] assert source.closed is True - assert mock_client.capture.call_count == 1 - assert mock_client.capture.call_args.kwargs["distinct_id"] == "test-user" + assert mock_client.capture_ai.call_count == 1 + assert mock_client.capture_ai.call_args.kwargs["distinct_id"] == "test-user" assert "posthog_distinct_id" not in client.post.call_args.kwargs @@ -2106,8 +2106,8 @@ async def test_async_messages_stream_preserves_provider_contract_and_manager( assert text == ["Hi"] assert source.closed is True - assert mock_client.capture.call_count == 1 - assert mock_client.capture.call_args.kwargs["distinct_id"] == "test-user" + assert mock_client.capture_ai.call_count == 1 + assert mock_client.capture_ai.call_args.kwargs["distinct_id"] == "test-user" assert "posthog_distinct_id" not in client.post.call_args.kwargs @@ -2153,7 +2153,7 @@ async def mock_async_create(**kwargs): events = [event async for event in stream] assert len(events) == 3 - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 def test_messages_streaming_early_exit_closes_provider_stream(mock_client): @@ -2176,7 +2176,7 @@ def test_messages_streaming_early_exit_closes_provider_stream(mock_client): break assert source.closed is True - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 @pytest.mark.asyncio @@ -2205,24 +2205,7 @@ async def mock_async_create(**kwargs): break assert source.closed is True - assert mock_client.capture.call_count == 1 - - -def test_ai_lane_client_routes_through_capture_ai(mock_client, mock_anthropic_response): - mock_client.enable_full_ai_capture = True - with patch( - "anthropic.resources.Messages.create", return_value=mock_anthropic_response - ): - client = Anthropic(api_key="test-key", posthog_client=mock_client) - client.messages.create( - model="claude-3-opus-20240229", - messages=[{"role": "user", "content": "Hello"}], - posthog_distinct_id="test-id", - ) - - mock_client.capture.assert_not_called() assert mock_client.capture_ai.call_count == 1 - assert mock_client.capture_ai.call_args[1]["event"] == "$ai_generation" def test_multimodal_client_skips_media_redaction(mock_client, mock_anthropic_response): diff --git a/posthog/test/ai/claude_agent_sdk/test_processor.py b/posthog/test/ai/claude_agent_sdk/test_processor.py index 1fb49c66..2b888b12 100644 --- a/posthog/test/ai/claude_agent_sdk/test_processor.py +++ b/posthog/test/ai/claude_agent_sdk/test_processor.py @@ -202,7 +202,7 @@ async def test_emits_generation_from_stream_events(self, processor, mock_client) collected.append(msg) # Should have captured $ai_generation + $ai_trace - calls = mock_client.capture.call_args_list + calls = mock_client.capture_ai.call_args_list events = [c.kwargs.get("event") or c[1].get("event") for c in calls] assert "$ai_generation" in events assert "$ai_trace" in events @@ -247,7 +247,7 @@ async def test_emits_cache_creation_ttl_breakdown_from_stream_events( generation_call = next( call - for call in mock_client.capture.call_args_list + for call in mock_client.capture_ai.call_args_list if (call.kwargs.get("event") or call[1].get("event")) == "$ai_generation" ) props = generation_call.kwargs.get("properties") or generation_call[1].get( @@ -286,7 +286,7 @@ async def test_emits_multiple_generations_for_multi_turn( gen_calls = [ c - for c in mock_client.capture.call_args_list + for c in mock_client.capture_ai.call_args_list if (c.kwargs.get("event") or c[1].get("event")) == "$ai_generation" ] assert len(gen_calls) == 2 @@ -318,7 +318,7 @@ async def test_fallback_generation_from_result_when_no_stream_events( gen_calls = [ c - for c in mock_client.capture.call_args_list + for c in mock_client.capture_ai.call_args_list if (c.kwargs.get("event") or c[1].get("event")) == "$ai_generation" ] assert len(gen_calls) == 1 @@ -349,7 +349,7 @@ async def test_fallback_copies_result_usage_before_capture( generation_call = next( call - for call in mock_client.capture.call_args_list + for call in mock_client.capture_ai.call_args_list if (call.kwargs.get("event") or call[1].get("event")) == "$ai_generation" ) props = generation_call.kwargs.get("properties") or generation_call[1].get( @@ -391,7 +391,7 @@ async def test_emits_span_for_tool_use(self, processor, mock_client): span_calls = [ c - for c in mock_client.capture.call_args_list + for c in mock_client.capture_ai.call_args_list if (c.kwargs.get("event") or c[1].get("event")) == "$ai_span" ] assert len(span_calls) == 1 @@ -422,7 +422,7 @@ async def test_emits_trace_on_result(self, processor, mock_client): trace_calls = [ c - for c in mock_client.capture.call_args_list + for c in mock_client.capture_ai.call_args_list if (c.kwargs.get("event") or c[1].get("event")) == "$ai_trace" ] assert len(trace_calls) == 1 @@ -449,7 +449,7 @@ async def test_trace_emits_error_status(self, processor, mock_client): trace_calls = [ c - for c in mock_client.capture.call_args_list + for c in mock_client.capture_ai.call_args_list if (c.kwargs.get("event") or c[1].get("event")) == "$ai_trace" ] assert len(trace_calls) == 1 @@ -473,7 +473,7 @@ async def test_client_privacy_mode_redacts_generation_input(self, mock_client): async for _ in proc.query(prompt="secret", options=ClaudeAgentOptions()): pass - properties = mock_client.capture.call_args.kwargs["properties"] + properties = mock_client.capture_ai.call_args.kwargs["properties"] assert properties["$ai_input"] is None @pytest.mark.asyncio @@ -502,7 +502,7 @@ async def test_per_call_privacy_mode_cannot_disable_processor_mode( ): pass - properties = mock_client.capture.call_args.kwargs["properties"] + properties = mock_client.capture_ai.call_args.kwargs["properties"] assert properties["$ai_input"] is None @pytest.mark.asyncio @@ -559,7 +559,7 @@ async def test_privacy_mode_redacts_tool_input(self, mock_client): span_calls = [ c - for c in mock_client.capture.call_args_list + for c in mock_client.capture_ai.call_args_list if (c.kwargs.get("event") or c[1].get("event")) == "$ai_span" ] assert len(span_calls) == 1 @@ -570,7 +570,7 @@ async def test_privacy_mode_redacts_tool_input(self, mock_client): generation_call = next( call - for call in mock_client.capture.call_args_list + for call in mock_client.capture_ai.call_args_list if (call.kwargs.get("event") or call[1].get("event")) == "$ai_generation" ) generation_props = generation_call.kwargs.get("properties") or generation_call[ @@ -619,7 +619,7 @@ async def test_data_url_prompt_and_tool_use_base64_are_redacted( gen_calls = [ c - for c in mock_client.capture.call_args_list + for c in mock_client.capture_ai.call_args_list if (c.kwargs.get("event") or c[1].get("event")) == "$ai_generation" ] assert len(gen_calls) == 2 @@ -667,7 +667,7 @@ async def test_no_distinct_id_sets_personless_option(self, mock_client): ): pass - for call in mock_client.capture.call_args_list: + for call in mock_client.capture_ai.call_args_list: assert call.kwargs["options"] == {"process_person_profile": False} assert call.kwargs["distinct_id"] == "trace-fallback" @@ -681,7 +681,7 @@ def test_explicit_groups_override_preserves_empty_groups(self, mock_client): proc._capture_event("$ai_trace", {}, groups={}) - assert mock_client.capture.call_args.kwargs["groups"] == {} + assert mock_client.capture_ai.call_args.kwargs["groups"] == {} def test_default_properties_keep_existing_precedence(self, mock_client): proc = PostHogClaudeAgentProcessor( @@ -694,7 +694,7 @@ def test_default_properties_keep_existing_precedence(self, mock_client): {"environment": "event", "$ai_trace_id": "trace-id"}, ) - assert mock_client.capture.call_args.kwargs["properties"] == { + assert mock_client.capture_ai.call_args.kwargs["properties"] == { "$ai_lib": "posthog-ai", "$ai_lib_version": VERSION, "environment": "processor", @@ -707,7 +707,7 @@ def test_client_without_capture_capability_is_ignored(self): proc._capture_event("$ai_trace", {}) def test_capture_errors_are_logged_and_suppressed(self, mock_client, caplog): - mock_client.capture.side_effect = RuntimeError("capture failed") + mock_client.capture_ai.side_effect = RuntimeError("capture failed") proc = PostHogClaudeAgentProcessor(client=mock_client) with caplog.at_level(logging.DEBUG, logger="posthog"): @@ -738,7 +738,7 @@ async def test_instance_properties_merged(self, mock_client): async for _ in proc.query(prompt="Hi", options=ClaudeAgentOptions()): pass - for call in mock_client.capture.call_args_list: + for call in mock_client.capture_ai.call_args_list: props = call.kwargs.get("properties") or call[1].get("properties") assert props.get("app") == "stamphog" assert props.get("version") == "1.0" @@ -763,7 +763,7 @@ async def test_per_call_properties_merged(self, processor, mock_client): ): pass - for call in mock_client.capture.call_args_list: + for call in mock_client.capture_ai.call_args_list: props = call.kwargs.get("properties") or call[1].get("properties") assert props.get("pr_number") == 42 @@ -791,7 +791,7 @@ def resolver(result): trace_calls = [ c - for c in mock_client.capture.call_args_list + for c in mock_client.capture_ai.call_args_list if (c.kwargs.get("event") or c[1].get("event")) == "$ai_trace" ] assert len(trace_calls) == 1 @@ -894,13 +894,3 @@ async def test_non_config_errors_propagate(self): prompt="Hello", options=ClaudeAgentOptions() ): pass - - -def test_ai_lane_client_routes_through_capture_ai(mock_client): - mock_client.enable_full_ai_capture = True - processor = PostHogClaudeAgentProcessor(client=mock_client, distinct_id="test-user") - processor._capture_event(event="$ai_trace", properties={}, distinct_id="d") - - mock_client.capture.assert_not_called() - mock_client.capture_ai.assert_called_once() - assert mock_client.capture_ai.call_args[1]["event"] == "$ai_trace" diff --git a/posthog/test/ai/gemini/test_gemini.py b/posthog/test/ai/gemini/test_gemini.py index 0079e73d..48d5cbd0 100644 --- a/posthog/test/ai/gemini/test_gemini.py +++ b/posthog/test/ai/gemini/test_gemini.py @@ -188,9 +188,9 @@ def test_new_client_basic_generation( ) assert response == mock_gemini_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -261,8 +261,8 @@ def mock_streaming_response(): 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] + assert mock_client.capture_ai.call_count == 1 + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -343,8 +343,8 @@ def mock_streaming_response(): 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] + assert mock_client.capture_ai.call_count == 1 + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -373,7 +373,7 @@ def test_new_client_groups(mock_client, mock_google_genai_client, mock_gemini_re posthog_groups={"company": "company_123"}, ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] assert call_args["groups"] == {"company": "company_123"} @@ -392,7 +392,7 @@ def test_new_client_privacy_mode_local( posthog_privacy_mode=True, ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] is None assert props["$ai_output_choices"] is None @@ -414,7 +414,7 @@ def test_new_client_privacy_mode_global( posthog_distinct_id="test-id", ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] is None assert props["$ai_output_choices"] is None @@ -432,7 +432,7 @@ def test_new_client_different_input_formats( client.models.generate_content( model="gemini-2.0-flash", contents="Hello", posthog_distinct_id="test-id" ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] == [{"role": "user", "content": "Hello"}] @@ -443,7 +443,7 @@ def test_new_client_different_input_formats( contents=[{"role": "user", "parts": [{"text": "hey"}]}], posthog_distinct_id="test-id", ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] == [ {"role": "user", "content": [{"type": "text", "text": "hey"}]} @@ -456,7 +456,7 @@ def test_new_client_different_input_formats( contents=[{"role": "user", "parts": [{"text": "Hello "}, {"text": "world"}]}], posthog_distinct_id="test-id", ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] == [ { @@ -469,11 +469,11 @@ def test_new_client_different_input_formats( ] # Test list input with string - mock_client.capture.reset_mock() + mock_client.capture_ai.reset_mock() client.models.generate_content( model="gemini-2.0-flash", contents=["List item"], posthog_distinct_id="test-id" ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] == [{"role": "user", "content": "List item"}] @@ -494,7 +494,7 @@ def test_new_client_model_parameters( max_tokens=100, ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_model_parameters"]["temperature"] == 0.7 assert props["$ai_model_parameters"]["max_tokens"] == 100 @@ -518,7 +518,7 @@ def test_new_client_default_settings( # Call without overriding defaults client.models.generate_content(model="gemini-2.0-flash", contents=["Hello"]) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "default_user" @@ -551,7 +551,7 @@ def test_new_client_override_defaults( posthog_groups={"organization": "special_org"}, ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Check overrides @@ -676,9 +676,9 @@ def test_tool_use_response(mock_client, mock_google_genai_client, mock_gemini_re ) assert response == mock_gemini_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -718,9 +718,9 @@ def test_function_calls_in_output_choices( ) assert response == mock_gemini_response_with_function_calls - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -767,9 +767,9 @@ def test_function_calls_only_no_content( ) assert response == mock_gemini_response_function_calls_only - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -826,9 +826,9 @@ def test_cache_and_reasoning_tokens(mock_client, mock_google_genai_client): ) assert response == mock_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Check that all token types are present @@ -890,9 +890,9 @@ def test_streaming_cache_and_reasoning_tokens(mock_client, mock_google_genai_cli assert len(result) == 2 # Check PostHog capture was called - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Check that all token types are present (should use final chunk's usage) @@ -963,9 +963,9 @@ def test_web_search_grounding(mock_client, mock_google_genai_client): ) assert response == mock_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify web search count is detected (binary for grounding) @@ -1034,9 +1034,9 @@ def mock_streaming_response(): chunks = list(response) assert len(chunks) == 2 - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify web search count is detected (binary for grounding) @@ -1098,9 +1098,9 @@ def test_empty_grounding_metadata_no_web_search(mock_client, mock_google_genai_c ) assert response == mock_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify web search count is 0 (not present in properties when 0) @@ -1161,9 +1161,9 @@ def test_empty_array_grounding_metadata_no_web_search( ) assert response == mock_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify web search count is 0 (not present in properties when 0) @@ -1203,9 +1203,9 @@ def test_falsy_candidates_no_web_search( ) assert response == mock_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Should not crash and web search count should not be present @@ -1257,9 +1257,9 @@ def test_embed_content_basic( ) assert response == mock_embed_content_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -1293,7 +1293,7 @@ def test_embed_content_with_token_counts( posthog_distinct_id="test-id", ) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_input_tokens"] == 13 # 5 + 8 @@ -1311,7 +1311,7 @@ def test_embed_content_without_token_counts( posthog_distinct_id="test-id", ) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] # No embedding carried a token count, so the property is omitted, not 0. assert "$ai_input_tokens" not in props @@ -1334,7 +1334,7 @@ def test_streaming_without_usage_omits_token_counts( ) list(response) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] assert "$ai_input_tokens" not in props assert "$ai_output_tokens" not in props @@ -1354,7 +1354,7 @@ def test_embed_content_privacy_mode( posthog_privacy_mode=True, ) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_input"] is None @@ -1371,7 +1371,7 @@ def test_embed_content_no_distinct_id( contents="Hello", ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Should fall back to trace_id as distinct_id @@ -1400,7 +1400,7 @@ def test_embed_content_default_params( posthog_properties={"extra": "prop"}, ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "default-id" @@ -1422,8 +1422,8 @@ def test_embed_content_error_handling(mock_client, mock_google_genai_client): ) # Event should still be captured - assert mock_client.capture.call_count == 1 - props = mock_client.capture.call_args[1]["properties"] + assert mock_client.capture_ai.call_count == 1 + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_is_error"] is True assert props["$ai_error"] == "API error" assert props["$ai_http_status"] == 0 @@ -1467,8 +1467,8 @@ def test_embed_content_integration(mock_client): assert all(isinstance(v, float) for v in response.embeddings[0].values) # Verify event captured - assert mock_client.capture.call_count == 1 - call_args = mock_client.capture.call_args[1] + assert mock_client.capture_ai.call_count == 1 + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["event"] == "$ai_embedding" assert props["$ai_provider"] == "gemini" @@ -1531,7 +1531,7 @@ def test_stop_reason_captured( posthog_distinct_id="test-id", ) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_stop_reason"] == finish_reason_name @@ -1580,8 +1580,8 @@ def mock_streaming_response(): list(response) - assert mock_client.capture.call_count == 1 - props = mock_client.capture.call_args[1]["properties"] + assert mock_client.capture_ai.call_count == 1 + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_stop_reason"] == "STOP" @@ -1597,26 +1597,8 @@ def test_integration_stop_reason(mock_client): posthog_distinct_id="test-id", ) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_stop_reason"] is not None assert isinstance(props["$ai_stop_reason"], str) assert props["$ai_provider"] == "gemini" assert props["$ai_input_tokens"] > 0 - - -def test_ai_lane_client_routes_through_capture_ai( - mock_client, mock_google_genai_client, mock_gemini_response -): - mock_google_genai_client.models.generate_content.return_value = mock_gemini_response - - mock_client.enable_full_ai_capture = True - client = Client(api_key="test-key", posthog_client=mock_client) - client.models.generate_content( - model="gemini-2.0-flash", - contents=["Tell me a fun fact about hedgehogs"], - posthog_distinct_id="test-id", - ) - - mock_client.capture.assert_not_called() - assert mock_client.capture_ai.call_count == 1 - assert mock_client.capture_ai.call_args[1]["event"] == "$ai_generation" diff --git a/posthog/test/ai/gemini/test_gemini_async.py b/posthog/test/ai/gemini/test_gemini_async.py index 1373a331..701b7ebd 100644 --- a/posthog/test/ai/gemini/test_gemini_async.py +++ b/posthog/test/ai/gemini/test_gemini_async.py @@ -134,9 +134,9 @@ async def test_async_client_basic_generation( ) assert response == mock_gemini_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -199,8 +199,8 @@ async def mock_streaming_response(): 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] + assert mock_client.capture_ai.call_count == 1 + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -283,8 +283,8 @@ async def mock_streaming_response(): 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] + assert mock_client.capture_ai.call_count == 1 + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -317,7 +317,7 @@ async def test_async_client_groups( posthog_groups={"company": "company_123"}, ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] assert call_args["groups"] == {"company": "company_123"} @@ -338,7 +338,7 @@ async def test_async_client_privacy_mode_local( posthog_privacy_mode=True, ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] is None assert props["$ai_output_choices"] is None @@ -362,7 +362,7 @@ async def test_async_client_privacy_mode_global( posthog_distinct_id="test-id", ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] is None assert props["$ai_output_choices"] is None @@ -382,7 +382,7 @@ async def test_async_client_different_input_formats( await client.models.generate_content( model="gemini-2.0-flash", contents="Hello", posthog_distinct_id="test-id" ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] == [{"role": "user", "content": "Hello"}] @@ -393,7 +393,7 @@ async def test_async_client_different_input_formats( contents=[{"role": "user", "parts": [{"text": "hey"}]}], posthog_distinct_id="test-id", ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] == [ {"role": "user", "content": [{"type": "text", "text": "hey"}]} @@ -406,7 +406,7 @@ async def test_async_client_different_input_formats( contents=[{"role": "user", "parts": [{"text": "Hello "}, {"text": "world"}]}], posthog_distinct_id="test-id", ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] == [ { @@ -419,11 +419,11 @@ async def test_async_client_different_input_formats( ] # Test list input with string - mock_client.capture.reset_mock() + mock_client.capture_ai.reset_mock() await client.models.generate_content( model="gemini-2.0-flash", contents=["List item"], posthog_distinct_id="test-id" ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] == [{"role": "user", "content": "List item"}] @@ -446,7 +446,7 @@ async def test_async_client_model_parameters( max_tokens=100, ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_model_parameters"]["temperature"] == 0.7 assert props["$ai_model_parameters"]["max_tokens"] == 100 @@ -472,7 +472,7 @@ async def test_async_client_default_settings( # Call without overriding defaults await client.models.generate_content(model="gemini-2.0-flash", contents=["Hello"]) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "default_user" @@ -507,7 +507,7 @@ async def test_async_client_override_defaults( posthog_groups={"organization": "special_org"}, ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Check overrides @@ -586,9 +586,9 @@ async def test_async_function_calls_in_output_choices( ) assert response == mock_gemini_response_with_function_calls - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -649,9 +649,9 @@ async def test_async_cache_and_reasoning_tokens(mock_client, mock_google_genai_c ) assert response == mock_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Check that all token types are present @@ -708,9 +708,9 @@ async def mock_streaming_response(): assert len(result) == 2 # Check PostHog capture was called - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Check that all token types are present (should use final chunk's usage) @@ -772,9 +772,9 @@ async def test_async_web_search_grounding(mock_client, mock_google_genai_client) ) assert response == mock_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify web search count is detected (binary for grounding) @@ -845,9 +845,9 @@ async def mock_streaming_response(): chunks.append(chunk) assert len(chunks) == 2 - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify web search count is detected (binary for grounding) @@ -899,9 +899,9 @@ async def test_async_embed_content_basic( ) assert response == mock_embed_content_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -935,7 +935,7 @@ async def test_async_embed_content_with_token_counts( posthog_distinct_id="test-id", ) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_input_tokens"] == 13 # 5 + 8 @@ -953,7 +953,7 @@ async def test_async_embed_content_without_token_counts( posthog_distinct_id="test-id", ) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] # No embedding carried a token count, so the property is omitted, not 0. assert "$ai_input_tokens" not in props @@ -973,7 +973,7 @@ async def test_async_embed_content_privacy_mode( posthog_privacy_mode=True, ) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_input"] is None @@ -990,7 +990,7 @@ async def test_async_embed_content_no_distinct_id( contents="Hello", ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Should fall back to trace_id as distinct_id @@ -1019,7 +1019,7 @@ async def test_async_embed_content_default_params( posthog_properties={"extra": "prop"}, ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "default-id" @@ -1045,8 +1045,8 @@ async def test_async_embed_content_error_handling( ) # Event should still be captured - assert mock_client.capture.call_count == 1 - props = mock_client.capture.call_args[1]["properties"] + assert mock_client.capture_ai.call_count == 1 + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_is_error"] is True assert props["$ai_error"] == "API error" assert props["$ai_http_status"] == 0 @@ -1090,8 +1090,8 @@ async def test_async_embed_content_integration(mock_client): assert all(isinstance(v, float) for v in response.embeddings[0].values) # Verify event captured - assert mock_client.capture.call_count == 1 - call_args = mock_client.capture.call_args[1] + assert mock_client.capture_ai.call_count == 1 + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["event"] == "$ai_embedding" assert props["$ai_provider"] == "gemini" @@ -1148,7 +1148,7 @@ async def mock_streaming_response(): chunks.append(chunk) assert len(chunks) == 1 - assert mock_client.capture.call_count == 1 - call_args = mock_client.capture.call_args[1] + assert mock_client.capture_ai.call_count == 1 + call_args = mock_client.capture_ai.call_args[1] assert call_args["event"] == "$ai_generation" assert call_args["properties"]["$ai_provider"] == "gemini" diff --git a/posthog/test/ai/gemini/test_gemini_parity.py b/posthog/test/ai/gemini/test_gemini_parity.py index de9f9b82..a3129903 100644 --- a/posthog/test/ai/gemini/test_gemini_parity.py +++ b/posthog/test/ai/gemini/test_gemini_parity.py @@ -178,8 +178,8 @@ async def test_sync_client_aio_models_tracks_generations( assert response is gemini_response provider_client.aio.models.generate_content.assert_awaited_once() - assert posthog_client.capture.call_count == 1 - call_args = posthog_client.capture.call_args[1] + assert posthog_client.capture_ai.call_count == 1 + call_args = posthog_client.capture_ai.call_args[1] assert call_args["distinct_id"] == "test-id" assert call_args["event"] == "$ai_generation" assert call_args["properties"]["$ai_model"] == "gemini-2.0-flash" diff --git a/posthog/test/ai/langchain/test_callbacks.py b/posthog/test/ai/langchain/test_callbacks.py index 7bc1b416..326eb0dd 100644 --- a/posthog/test/ai/langchain/test_callbacks.py +++ b/posthog/test/ai/langchain/test_callbacks.py @@ -185,15 +185,15 @@ def test_basic_chat_chain(mock_client, stream): result = chain.invoke({}, config={"callbacks": callbacks}) assert result.content == "The Los Angeles Dodgers won the World Series in 2020." - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - span_args = mock_client.capture.call_args_list[0][1] + span_args = mock_client.capture_ai.call_args_list[0][1] span_props = span_args["properties"] - generation_args = mock_client.capture.call_args_list[1][1] + generation_args = mock_client.capture_ai.call_args_list[1][1] generation_props = generation_args["properties"] - trace_args = mock_client.capture.call_args_list[2][1] + trace_args = mock_client.capture_ai.call_args_list[2][1] trace_props = trace_args["properties"] # Span is first @@ -262,13 +262,13 @@ async def test_async_basic_chat_chain(mock_client, stream): else: result = await chain.ainvoke({}, config={"callbacks": callbacks}) assert result.content == "The Los Angeles Dodgers won the World Series in 2020." - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - span_args = mock_client.capture.call_args_list[0][1] + span_args = mock_client.capture_ai.call_args_list[0][1] span_props = span_args["properties"] - generation_args = mock_client.capture.call_args_list[1][1] + generation_args = mock_client.capture_ai.call_args_list[1][1] generation_props = generation_args["properties"] - trace_args = mock_client.capture.call_args_list[2][1] + trace_args = mock_client.capture_ai.call_args_list[2][1] trace_props = trace_args["properties"] # Span is first @@ -349,7 +349,7 @@ def _capture_anthropic_usage( ) cb._pop_run_and_capture_generation(run_id, None, response) - return mock_client.capture.call_args.kwargs["properties"] + return mock_client.capture_ai.call_args.kwargs["properties"] @pytest.mark.parametrize( @@ -380,8 +380,8 @@ def test_basic_llm_chain(mock_client, Model, stream): ) assert result == "The Los Angeles Dodgers won the World Series in 2020." - assert mock_client.capture.call_count == 1 - args = mock_client.capture.call_args_list[0][1] + assert mock_client.capture_ai.call_count == 1 + args = mock_client.capture_ai.call_args_list[0][1] props = args["properties"] assert args["event"] == "$ai_generation" @@ -425,8 +425,8 @@ async def test_async_basic_llm_chain(mock_client, Model, stream): ) assert result == "The Los Angeles Dodgers won the World Series in 2020." - assert mock_client.capture.call_count == 1 - args = mock_client.capture.call_args_list[0][1] + assert mock_client.capture_ai.call_count == 1 + args = mock_client.capture_ai.call_args_list[0][1] props = args["properties"] assert args["event"] == "$ai_generation" @@ -455,21 +455,21 @@ def test_trace_id_and_inputs_for_multiple_chains(mock_client): assert result.content == "Bar" # span, generation, span, generation, trace - assert mock_client.capture.call_count == 5 + assert mock_client.capture_ai.call_count == 5 - first_span_args = mock_client.capture.call_args_list[0][1] + first_span_args = mock_client.capture_ai.call_args_list[0][1] first_span_props = first_span_args["properties"] - first_generation_args = mock_client.capture.call_args_list[1][1] + first_generation_args = mock_client.capture_ai.call_args_list[1][1] first_generation_props = first_generation_args["properties"] - second_span_args = mock_client.capture.call_args_list[2][1] + second_span_args = mock_client.capture_ai.call_args_list[2][1] second_span_props = second_span_args["properties"] - second_generation_args = mock_client.capture.call_args_list[3][1] + second_generation_args = mock_client.capture_ai.call_args_list[3][1] second_generation_props = second_generation_args["properties"] - trace_args = mock_client.capture.call_args_list[4][1] + trace_args = mock_client.capture_ai.call_args_list[4][1] trace_props = trace_args["properties"] # Prompt span @@ -536,10 +536,10 @@ def test_personless_mode(mock_client): prompt = ChatPromptTemplate.from_messages([("user", "Foo")]) chain = prompt | FakeMessagesListChatModel(responses=[AIMessage(content="Bar")]) chain.invoke({}, config={"callbacks": [CallbackHandler(mock_client)]}) - assert mock_client.capture.call_count == 3 - span_args = mock_client.capture.call_args_list[0][1] - generation_args = mock_client.capture.call_args_list[1][1] - trace_args = mock_client.capture.call_args_list[2][1] + assert mock_client.capture_ai.call_count == 3 + span_args = mock_client.capture_ai.call_args_list[0][1] + generation_args = mock_client.capture_ai.call_args_list[1][1] + trace_args = mock_client.capture_ai.call_args_list[2][1] # span assert span_args["event"] == "$ai_span" @@ -555,10 +555,10 @@ def test_personless_mode(mock_client): chain.invoke( {}, config={"callbacks": [CallbackHandler(mock_client, distinct_id=id)]} ) - assert mock_client.capture.call_count == 6 - span_args = mock_client.capture.call_args_list[3][1] - generation_args = mock_client.capture.call_args_list[4][1] - trace_args = mock_client.capture.call_args_list[5][1] + assert mock_client.capture_ai.call_count == 6 + span_args = mock_client.capture_ai.call_args_list[3][1] + generation_args = mock_client.capture_ai.call_args_list[4][1] + trace_args = mock_client.capture_ai.call_args_list[5][1] # span assert "options" not in span_args @@ -595,10 +595,10 @@ def test_personless_mode_exception(mock_client, unauthorized_http_client): callbacks = CallbackHandler(mock_client) with pytest.raises(AuthenticationError): chain.invoke({}, config={"callbacks": [callbacks]}) - assert mock_client.capture.call_count == 3 - span_args = mock_client.capture.call_args_list[0][1] - generation_args = mock_client.capture.call_args_list[1][1] - trace_args = mock_client.capture.call_args_list[2][1] + assert mock_client.capture_ai.call_count == 3 + span_args = mock_client.capture_ai.call_args_list[0][1] + generation_args = mock_client.capture_ai.call_args_list[1][1] + trace_args = mock_client.capture_ai.call_args_list[2][1] # span assert span_args["event"] == "$ai_span" @@ -615,10 +615,10 @@ def test_personless_mode_exception(mock_client, unauthorized_http_client): chain.invoke( {}, config={"callbacks": [CallbackHandler(mock_client, distinct_id=id)]} ) - assert mock_client.capture.call_count == 6 - span_args = mock_client.capture.call_args_list[3][1] - generation_args = mock_client.capture.call_args_list[4][1] - trace_args = mock_client.capture.call_args_list[5][1] + assert mock_client.capture_ai.call_count == 6 + span_args = mock_client.capture_ai.call_args_list[3][1] + generation_args = mock_client.capture_ai.call_args_list[4][1] + trace_args = mock_client.capture_ai.call_args_list[5][1] # span assert span_args["event"] == "$ai_span" @@ -655,9 +655,9 @@ def test_metadata(mock_client): result = chain.invoke({"plan": None}, config={"callbacks": callbacks}) assert result.content == "Bar" - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - span_call_args = mock_client.capture.call_args_list[0][1] + span_call_args = mock_client.capture_ai.call_args_list[0][1] span_call_props = span_call_args["properties"] assert span_call_args["distinct_id"] == "test_id" assert span_call_args["event"] == "$ai_span" @@ -666,7 +666,7 @@ def test_metadata(mock_client): assert "$ai_parent_id" in span_call_props assert "$ai_span_id" in span_call_props - generation_call_args = mock_client.capture.call_args_list[1][1] + generation_call_args = mock_client.capture_ai.call_args_list[1][1] generation_call_props = generation_call_args["properties"] assert generation_call_args["distinct_id"] == "test_id" assert generation_call_args["event"] == "$ai_generation" @@ -679,7 +679,7 @@ def test_metadata(mock_client): assert generation_call_props["$ai_http_status"] == 200 assert isinstance(generation_call_props["$ai_latency"], float) - trace_call_args = mock_client.capture.call_args_list[2][1] + trace_call_args = mock_client.capture_ai.call_args_list[2][1] trace_call_props = trace_call_args["properties"] assert trace_call_args["distinct_id"] == "test_id" assert trace_call_args["event"] == "$ai_trace" @@ -742,8 +742,8 @@ def test_graph_state(mock_client): assert isinstance(result["messages"][2], AIMessage) assert result["messages"][2].content == "It's a type of greeble." - assert mock_client.capture.call_count == 6 - calls = [call[1] for call in mock_client.capture.call_args_list] + assert mock_client.capture_ai.call_count == 6 + calls = [call[1] for call in mock_client.capture_ai.call_args_list] # The trace event is captured at the end trace_args = calls[-1] @@ -867,8 +867,8 @@ def runnable(_): assert callbacks._runs == {} assert callbacks._parent_tree == {} - assert mock_client.capture.call_count == 1 - trace_call_args = mock_client.capture.call_args_list[0][1] + assert mock_client.capture_ai.call_count == 1 + trace_call_args = mock_client.capture_ai.call_args_list[0][1] assert trace_call_args["event"] == "$ai_trace" assert trace_call_args["properties"]["$ai_span_name"] == "runnable" @@ -895,8 +895,8 @@ def unauthorized(request): assert callbacks._runs == {} assert callbacks._parent_tree == {} - assert mock_client.capture.call_count == 3 - generation_args = mock_client.capture.call_args_list[1][1] + assert mock_client.capture_ai.call_count == 3 + generation_args = mock_client.capture_ai.call_args_list[1][1] props = generation_args["properties"] assert props["$ai_http_status"] == 401 assert props["$ai_input"] == [{"role": "user", "content": "Foo"}] @@ -928,9 +928,9 @@ def test_openai_chain(mock_client): approximate_latency = math.floor(time.time() - start_time) assert result.content == "Bar" - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - gen_args = mock_client.capture.call_args_list[1][1] + gen_args = mock_client.capture_ai.call_args_list[1][1] gen_props = gen_args["properties"] assert gen_args["event"] == "$ai_generation" assert gen_props["$ai_trace_id"] == "test-trace-id" @@ -990,11 +990,11 @@ def test_openai_captures_multiple_generations(mock_client): result = chain.invoke({}, config={"callbacks": [callbacks]}) assert result.content == "Bar" - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - gen_args = mock_client.capture.call_args_list[1][1] + gen_args = mock_client.capture_ai.call_args_list[1][1] gen_props = gen_args["properties"] - trace_args = mock_client.capture.call_args_list[2][1] + trace_args = mock_client.capture_ai.call_args_list[2][1] trace_props = trace_args["properties"] assert gen_args["event"] == "$ai_generation" @@ -1053,11 +1053,11 @@ def test_openai_streaming(mock_client): result = sum(result[1:], result[0]) assert result.content == "Bar" - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - gen_args = mock_client.capture.call_args_list[1][1] + gen_args = mock_client.capture_ai.call_args_list[1][1] gen_props = gen_args["properties"] - trace_args = mock_client.capture.call_args_list[2][1] + trace_args = mock_client.capture_ai.call_args_list[2][1] trace_props = trace_args["properties"] assert gen_args["event"] == "$ai_generation" @@ -1097,11 +1097,11 @@ async def test_async_openai_streaming(mock_client): result = sum(result[1:], result[0]) assert result.content == "Bar" - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - gen_args = mock_client.capture.call_args_list[1][1] + gen_args = mock_client.capture_ai.call_args_list[1][1] gen_props = gen_args["properties"] - trace_args = mock_client.capture.call_args_list[2][1] + trace_args = mock_client.capture_ai.call_args_list[2][1] trace_props = trace_args["properties"] assert gen_args["event"] == "$ai_generation" @@ -1135,8 +1135,8 @@ def test_base_url_retrieval(mock_client, unauthorized_http_client): with pytest.raises(AuthenticationError): chain.invoke({}, config={"callbacks": [callbacks]}) - assert mock_client.capture.call_count == 3 - generation_call = mock_client.capture.call_args_list[1][1] + assert mock_client.capture_ai.call_count == 3 + generation_call = mock_client.capture_ai.call_args_list[1][1] assert generation_call["properties"]["$ai_base_url"] == "https://test.posthog.com" @@ -1152,8 +1152,8 @@ def test_groups(mock_client): callbacks = CallbackHandler(mock_client, groups={"company": "test_company"}) chain.invoke({}, config={"callbacks": [callbacks]}) - assert mock_client.capture.call_count == 3 - generation_call = mock_client.capture.call_args_list[1][1] + assert mock_client.capture_ai.call_count == 3 + generation_call = mock_client.capture_ai.call_args_list[1][1] assert generation_call["groups"] == {"company": "test_company"} @@ -1169,8 +1169,8 @@ def test_privacy_mode_local(mock_client): callbacks = CallbackHandler(mock_client, privacy_mode=True) chain.invoke({}, config={"callbacks": [callbacks]}) - assert mock_client.capture.call_count == 3 - generation_call = mock_client.capture.call_args_list[1][1] + assert mock_client.capture_ai.call_count == 3 + generation_call = mock_client.capture_ai.call_args_list[1][1] assert generation_call["properties"]["$ai_input"] is None assert generation_call["properties"]["$ai_output_choices"] is None @@ -1188,8 +1188,8 @@ def test_privacy_mode_global(mock_client): callbacks = CallbackHandler(mock_client) chain.invoke({}, config={"callbacks": [callbacks]}) - assert mock_client.capture.call_count == 3 - generation_call = mock_client.capture.call_args_list[1][1] + assert mock_client.capture_ai.call_count == 3 + generation_call = mock_client.capture_ai.call_args_list[1][1] assert generation_call["properties"]["$ai_input"] is None assert generation_call["properties"]["$ai_output_choices"] is None @@ -1219,11 +1219,11 @@ def test_anthropic_chain(mock_client): approximate_latency = math.floor(time.time() - start_time) assert result.content == "Bar" - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - gen_args = mock_client.capture.call_args_list[1][1] + gen_args = mock_client.capture_ai.call_args_list[1][1] gen_props = gen_args["properties"] - trace_args = mock_client.capture.call_args_list[2][1] + trace_args = mock_client.capture_ai.call_args_list[2][1] trace_props = trace_args["properties"] assert gen_args["event"] == "$ai_generation" @@ -1279,11 +1279,11 @@ async def test_async_anthropic_streaming(mock_client): result = sum(result[1:], result[0]) assert result.content == "Bar" - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - gen_args = mock_client.capture.call_args_list[1][1] + gen_args = mock_client.capture_ai.call_args_list[1][1] gen_props = gen_args["properties"] - trace_args = mock_client.capture.call_args_list[2][1] + trace_args = mock_client.capture_ai.call_args_list[2][1] trace_props = trace_args["properties"] assert gen_args["event"] == "$ai_generation" @@ -1385,8 +1385,8 @@ def test_tool_calls(mock_client): callbacks = CallbackHandler(mock_client) chain.invoke({}, config={"callbacks": [callbacks]}) - assert mock_client.capture.call_count == 3 - generation_call = mock_client.capture.call_args_list[1][1] + assert mock_client.capture_ai.call_count == 3 + generation_call = mock_client.capture_ai.call_args_list[1][1] assert generation_call["properties"]["$ai_output_choices"][0]["tool_calls"] == [ { "type": "function", @@ -1420,10 +1420,10 @@ async def sleep(x): # -> Any: chain2.ainvoke({}, config={"callbacks": [cb]}), ) approximate_latency = math.floor(time.time() - start_time) - assert mock_client.capture.call_count == 4 + assert mock_client.capture_ai.call_count == 4 first_call, second_call, third_call, fourth_call = ( - mock_client.capture.call_args_list + mock_client.capture_ai.call_args_list ) assert first_call[1]["event"] == "$ai_span" assert second_call[1]["event"] == "$ai_generation" @@ -1460,7 +1460,7 @@ def get_weather(city: Literal["nyc", "sf"]): mock_client, trace_id="test-trace-id", distinct_id="test-distinct-id" ) graph.invoke(inputs, config={"callbacks": [cb]}) - calls = [call[1] for call in mock_client.capture.call_args_list] + calls = [call[1] for call in mock_client.capture_ai.call_args_list] assert len(calls) == 15 for call in calls: assert call["properties"]["$ai_trace_id"] == "test-trace-id" @@ -1486,15 +1486,15 @@ def test_span_set_parent_ids(mock_client, trace_id): chain = prompt | model chain.invoke({}, config={"callbacks": callbacks}) - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - span_props = mock_client.capture.call_args_list[0][1] + span_props = mock_client.capture_ai.call_args_list[0][1] assert ( span_props["properties"]["$ai_trace_id"] == span_props["properties"]["$ai_parent_id"] ) - generation_props = mock_client.capture.call_args_list[1][1] + generation_props = mock_client.capture_ai.call_args_list[1][1] assert ( generation_props["properties"]["$ai_trace_id"] == generation_props["properties"]["$ai_parent_id"] @@ -1516,9 +1516,9 @@ def span_3(_): chain = RunnableLambda(span_1) chain.invoke({}, config={"callbacks": callbacks}) - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - calls = mock_client.capture.call_args_list + calls = mock_client.capture_ai.call_args_list span_props_2 = calls[0][1]["properties"] span_props_1 = calls[1][1]["properties"] trace_props = calls[2][1]["properties"] @@ -1538,12 +1538,12 @@ def span(_): except ValueError: pass - assert mock_client.capture.call_count == 2 + assert mock_client.capture_ai.call_count == 2 assert ( - mock_client.capture.call_args_list[1][1]["properties"]["$ai_error"] + mock_client.capture_ai.call_args_list[1][1]["properties"]["$ai_error"] == "ValueError: test" ) - assert mock_client.capture.call_args_list[1][1]["properties"]["$ai_is_error"] + assert mock_client.capture_ai.call_args_list[1][1]["properties"]["$ai_is_error"] def test_captures_error_without_details_in_span(mock_client): @@ -1557,12 +1557,12 @@ def span(_): except ValueError: pass - assert mock_client.capture.call_count == 2 + assert mock_client.capture_ai.call_count == 2 assert ( - mock_client.capture.call_args_list[1][1]["properties"]["$ai_error"] + mock_client.capture_ai.call_args_list[1][1]["properties"]["$ai_error"] == "ValueError" ) - assert mock_client.capture.call_args_list[1][1]["properties"]["$ai_is_error"] + assert mock_client.capture_ai.call_args_list[1][1]["properties"]["$ai_is_error"] def test_openai_reasoning_tokens(mock_client): @@ -1591,9 +1591,9 @@ def test_openai_reasoning_tokens(mock_client): result = chain.invoke({}, config={"callbacks": callbacks}) assert result.content == "Let me think through this step by step..." - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - generation_args = mock_client.capture.call_args_list[1][1] + generation_args = mock_client.capture_ai.call_args_list[1][1] generation_props = generation_args["properties"] assert generation_args["event"] == "$ai_generation" @@ -1626,9 +1626,9 @@ def test_anthropic_cache_write_and_read_tokens(mock_client): result = chain.invoke({}, config={"callbacks": callbacks}) assert result.content == "I've analyzed the document and cached the context." - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - generation_args = mock_client.capture.call_args_list[1][1] + generation_args = mock_client.capture_ai.call_args_list[1][1] generation_props = generation_args["properties"] assert generation_args["event"] == "$ai_generation" @@ -1661,9 +1661,9 @@ def test_anthropic_cache_write_and_read_tokens(mock_client): result = chain.invoke({}, config={"callbacks": callbacks}) assert result.content == "Using cached analysis to provide quick response." - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - generation_args = mock_client.capture.call_args_list[1][1] + generation_args = mock_client.capture_ai.call_args_list[1][1] generation_props = generation_args["properties"] assert generation_args["event"] == "$ai_generation" @@ -1714,7 +1714,7 @@ def test_anthropic_provider_subtracts_cache_tokens(mock_client): cb._pop_run_and_capture_generation(run_id, None, response) - generation_args = mock_client.capture.call_args_list[0][1] + generation_args = mock_client.capture_ai.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 @@ -1757,7 +1757,7 @@ def test_anthropic_provider_subtracts_cache_write_tokens(mock_client): cb._pop_run_and_capture_generation(run_id, None, response) - generation_args = mock_client.capture.call_args_list[0][1] + generation_args = mock_client.capture_ai.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 @@ -1930,7 +1930,7 @@ def test_anthropic_provider_subtracts_both_cache_read_and_write_tokens(mock_clie cb._pop_run_and_capture_generation(run_id, None, response) - generation_args = mock_client.capture.call_args_list[0][1] + generation_args = mock_client.capture_ai.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 @@ -1966,9 +1966,9 @@ def test_openai_cache_read_tokens(mock_client): result = chain.invoke({}, config={"callbacks": callbacks}) assert result.content == "Response using cached prompt context." - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - generation_args = mock_client.capture.call_args_list[1][1] + generation_args = mock_client.capture_ai.call_args_list[1][1] generation_props = generation_args["properties"] assert generation_args["event"] == "$ai_generation" @@ -2010,9 +2010,9 @@ def test_openai_cache_creation_tokens(mock_client): result = chain.invoke({}, config={"callbacks": callbacks}) assert result.content == "Created cache for the prompt context." - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - generation_args = mock_client.capture.call_args_list[1][1] + generation_args = mock_client.capture_ai.call_args_list[1][1] generation_props = generation_args["properties"] assert generation_args["event"] == "$ai_generation" @@ -2052,9 +2052,9 @@ def test_combined_reasoning_and_cache_tokens(mock_client): result = chain.invoke({}, config={"callbacks": callbacks}) assert result.content == "Let me reason through this using cached context..." - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - generation_args = mock_client.capture.call_args_list[1][1] + generation_args = mock_client.capture_ai.call_args_list[1][1] generation_props = generation_args["properties"] assert generation_args["event"] == "$ai_generation" @@ -2074,7 +2074,7 @@ def test_openai_reasoning_tokens_o4_mini(mock_client): mock_client, trace_id="test-trace-id", distinct_id="test-distinct-id" ) model.invoke("what is the weather in sf", config={"callbacks": [cb]}) - call = mock_client.capture.call_args_list[0][1] + call = mock_client.capture_ai.call_args_list[0][1] assert call["properties"]["$ai_reasoning_tokens"] is not None assert call["properties"]["$ai_input_tokens"] is not None assert call["properties"]["$ai_output_tokens"] is not None @@ -2103,7 +2103,7 @@ def test_callback_handler_without_client(): assert result.content == "Bar" # Verify that the mock client was used for capturing events - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 def test_convert_message_to_dict_tool_calls(): @@ -2210,8 +2210,8 @@ def test_tool_definition(mock_client): parent_run_id=None, ) - assert mock_client.capture.call_count == 1 - call_args = mock_client.capture.call_args[1] + assert mock_client.capture_ai.call_count == 1 + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == run_id @@ -2261,9 +2261,9 @@ def test_cache_read_tokens_subtraction_from_input_tokens(mock_client): result = chain.invoke({}, config={"callbacks": callbacks}) assert result.content == "Response using cached prompt context." - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - generation_args = mock_client.capture.call_args_list[1][1] + generation_args = mock_client.capture_ai.call_args_list[1][1] generation_props = generation_args["properties"] assert generation_args["event"] == "$ai_generation" @@ -2303,9 +2303,9 @@ def test_cache_read_tokens_subtraction_prevents_negative(mock_client): result = chain.invoke({}, config={"callbacks": callbacks}) assert result.content == "Response with edge case token counts." - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - generation_args = mock_client.capture.call_args_list[1][1] + generation_args = mock_client.capture_ai.call_args_list[1][1] generation_props = generation_args["properties"] assert generation_args["event"] == "$ai_generation" @@ -2344,9 +2344,9 @@ def test_no_cache_read_tokens_no_subtraction(mock_client): result = chain.invoke({}, config={"callbacks": callbacks}) assert result.content == "Response without cache." - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - generation_args = mock_client.capture.call_args_list[1][1] + generation_args = mock_client.capture_ai.call_args_list[1][1] generation_props = generation_args["properties"] assert generation_args["event"] == "$ai_generation" @@ -2383,9 +2383,9 @@ def test_zero_input_tokens_with_cache_read(mock_client): result = chain.invoke({}, config={"callbacks": callbacks}) assert result.content == "Response." - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - generation_args = mock_client.capture.call_args_list[1][1] + generation_args = mock_client.capture_ai.call_args_list[1][1] generation_props = generation_args["properties"] assert generation_args["event"] == "$ai_generation" @@ -2424,9 +2424,9 @@ def test_non_anthropic_cache_write_tokens_not_subtracted_from_input(mock_client) result = chain.invoke({}, config={"callbacks": callbacks}) assert result.content == "Creating cache." - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - generation_args = mock_client.capture.call_args_list[1][1] + generation_args = mock_client.capture_ai.call_args_list[1][1] generation_props = generation_args["properties"] assert generation_args["event"] == "$ai_generation" @@ -2475,8 +2475,8 @@ def test_agent_action_and_finish_imports(): callbacks.on_agent_finish(finish, run_id=run_id, parent_run_id=parent_run_id) # Verify capture was called - assert mock_client.capture.call_count == 1 - call_args = mock_client.capture.call_args[1] + assert mock_client.capture_ai.call_count == 1 + call_args = mock_client.capture_ai.call_args[1] assert call_args["event"] == "$ai_span" @@ -2583,8 +2583,8 @@ def test_billable_property_in_generation_event(mock_client): parent_run_id=None, ) - assert mock_client.capture.call_count == 1 - call_args = mock_client.capture.call_args[1] + assert mock_client.capture_ai.call_count == 1 + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["event"] == "$ai_generation" @@ -2603,7 +2603,7 @@ def test_billable_defaults_to_false_in_event(mock_client): chain.invoke({}, config={"callbacks": callbacks}) generation_call = None - for call in mock_client.capture.call_args_list: + for call in mock_client.capture_ai.call_args_list: if call[1]["event"] == "$ai_generation": generation_call = call break @@ -2647,8 +2647,8 @@ def test_billable_with_real_chain(mock_client): parent_run_id=None, ) - assert mock_client.capture.call_count == 1 - call_args = mock_client.capture.call_args[1] + assert mock_client.capture_ai.call_count == 1 + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["event"] == "$ai_generation" @@ -2706,7 +2706,7 @@ def failing_span(_): # Find the span event (should have $ai_is_error=True) span_calls = [ call - for call in mock_client.capture.call_args_list + for call in mock_client.capture_ai.call_args_list if call[1].get("properties", {}).get("$ai_is_error") is True ] assert len(span_calls) >= 1 @@ -2739,7 +2739,7 @@ def failing_span(_): # But the span event should still have error info span_calls = [ call - for call in mock_client.capture.call_args_list + for call in mock_client.capture_ai.call_args_list if call[1].get("properties", {}).get("$ai_is_error") is True ] assert len(span_calls) >= 1 @@ -2777,7 +2777,7 @@ def test_exception_autocapture_on_llm_generation_error(mock_client): # Verify the generation event has $exception_event_id generation_calls = [ call - for call in mock_client.capture.call_args_list + for call in mock_client.capture_ai.call_args_list if call[1].get("event") == "$ai_generation" ] assert len(generation_calls) == 1 @@ -2851,7 +2851,7 @@ def failing_span(_): # Span event should NOT have $exception_event_id span_calls = [ call - for call in mock_client.capture.call_args_list + for call in mock_client.capture_ai.call_args_list if call[1].get("properties", {}).get("$ai_is_error") is True ] assert len(span_calls) >= 1 @@ -2860,20 +2860,6 @@ def failing_span(_): assert "$exception_event_id" not in span_props -def test_ai_lane_client_routes_through_capture_ai(mock_client): - prompt = ChatPromptTemplate.from_messages([("user", "Who won the world series?")]) - model = FakeMessagesListChatModel(responses=[AIMessage(content="The Dodgers.")]) - mock_client.enable_full_ai_capture = True - callbacks = [CallbackHandler(mock_client)] - - (prompt | model).invoke({}, config={"callbacks": callbacks}) - - mock_client.capture.assert_not_called() - events = [c[1]["event"] for c in mock_client.capture_ai.call_args_list] - assert "$ai_generation" in events - assert "$ai_trace" in events - - def test_served_service_tier_merges_into_model_parameters(mock_client): from langchain_core.messages import AIMessage from langchain_core.outputs import ChatGeneration, LLMResult @@ -2894,7 +2880,7 @@ def test_served_service_tier_merges_into_model_parameters(mock_client): cb._pop_run_and_capture_generation(run_id, None, response) - props = mock_client.capture.call_args.kwargs["properties"] + props = mock_client.capture_ai.call_args.kwargs["properties"] assert props["$ai_model_parameters"]["service_tier"] == "flex" assert props["$ai_model_parameters"]["temperature"] == 0.5 assert props["$ai_service_tier"] == "flex" @@ -2951,5 +2937,5 @@ def test_stop_reason_resolution( cb._pop_run_and_capture_generation(run_id, None, response) - props = mock_client.capture.call_args.kwargs["properties"] + props = mock_client.capture_ai.call_args.kwargs["properties"] assert props.get("$ai_stop_reason") == expected diff --git a/posthog/test/ai/langchain/test_middleware.py b/posthog/test/ai/langchain/test_middleware.py index ed68bca8..441cdc22 100644 --- a/posthog/test/ai/langchain/test_middleware.py +++ b/posthog/test/ai/langchain/test_middleware.py @@ -121,7 +121,7 @@ def _client() -> MagicMock: def _events(client: MagicMock) -> list[dict[str, Any]]: - return [call.kwargs for call in client.capture.call_args_list] + return [call.kwargs for call in client.capture_ai.call_args_list] def _event(client: MagicMock, event_name: str) -> dict[str, Any]: @@ -353,7 +353,7 @@ async def test_async_hooks_do_not_capture_on_the_event_loop_thread() -> None: client.sync_mode = True event_loop_thread = threading.get_ident() capture_threads: list[int] = [] - client.capture.side_effect = lambda **_: capture_threads.append( + client.capture_ai.side_effect = lambda **_: capture_threads.append( threading.get_ident() ) middleware = PostHogMiddleware(client) @@ -841,7 +841,7 @@ def test_privacy_mode_redacts_agent_and_model_content() -> None: def test_capture_failure_never_changes_agent_result() -> None: client = _client() - client.capture.side_effect = RuntimeError("telemetry unavailable") + client.capture_ai.side_effect = RuntimeError("telemetry unavailable") middleware = PostHogMiddleware(client) state = _middleware_state(middleware) request = _model_request(StubAgentModel(responses=[]), state) diff --git a/posthog/test/ai/openai/test_async_parity.py b/posthog/test/ai/openai/test_async_parity.py index 75ab907c..8a96bfec 100644 --- a/posthog/test/ai/openai/test_async_parity.py +++ b/posthog/test/ai/openai/test_async_parity.py @@ -46,7 +46,7 @@ def _sync_props(mock_client, chunks): posthog_distinct_id="test-id", ) ) - return mock_client.capture.call_args[1]["properties"] + return mock_client.capture_ai.call_args[1]["properties"] async def _async_props(mock_client, chunks): @@ -68,7 +68,7 @@ async def it(): ) async for _ in stream: pass - return mock_client.capture.call_args[1]["properties"] + return mock_client.capture_ai.call_args[1]["properties"] @pytest.mark.asyncio @@ -76,7 +76,7 @@ async def test_async_streaming_emits_the_same_properties_as_sync( mock_client, streaming_tool_call_chunks ): sync_props = _sync_props(mock_client, streaming_tool_call_chunks) - mock_client.capture.reset_mock() + mock_client.capture_ai.reset_mock() async_props = await _async_props(mock_client, streaming_tool_call_chunks) # Guard: if the sync side stopped emitting these, the comparison below is vacuous. @@ -120,7 +120,7 @@ async def test_responses_streaming_properties_have_sync_async_parity(mock_client ): client = OpenAI(api_key="test-key", posthog_client=mock_client) list(client.responses.create(**request)) - sync_props = mock_client.capture.call_args.kwargs["properties"] + sync_props = mock_client.capture_ai.call_args.kwargs["properties"] async def create(self, **kwargs): async def chunks(): @@ -128,13 +128,13 @@ async def chunks(): return chunks() - mock_client.capture.reset_mock() + mock_client.capture_ai.reset_mock() with patch("openai.resources.responses.AsyncResponses.create", new=create): client = AsyncOpenAI(api_key="test-key", posthog_client=mock_client) stream = await client.responses.create(**request) async for _ in stream: pass - async_props = mock_client.capture.call_args.kwargs["properties"] + async_props = mock_client.capture_ai.call_args.kwargs["properties"] sync_without_latency = {k: v for k, v in sync_props.items() if k != "$ai_latency"} async_without_latency = {k: v for k, v in async_props.items() if k != "$ai_latency"} @@ -167,15 +167,15 @@ async def test_embedding_telemetry_has_sync_async_parity(mock_client): client = OpenAI(api_key="test-key", posthog_client=mock_client) assert client.embeddings.create(**request) is response sync_create.assert_called_once_with(**provider_request) - sync_capture = mock_client.capture.call_args + sync_capture = mock_client.capture_ai.call_args async_create = AsyncMock(return_value=response) - mock_client.capture.reset_mock() + mock_client.capture_ai.reset_mock() with patch("openai.resources.embeddings.AsyncEmbeddings.create", new=async_create): client = AsyncOpenAI(api_key="test-key", posthog_client=mock_client) assert await client.embeddings.create(**request) is response async_create.assert_awaited_once_with(**provider_request) - async_capture = mock_client.capture.call_args + async_capture = mock_client.capture_ai.call_args sync_props = sync_capture.kwargs["properties"] async_props = async_capture.kwargs["properties"] @@ -209,5 +209,5 @@ def test_sync_stream_close_after_early_exit_captures_partial_state( assert next(stream) == streaming_tool_call_chunks[0] stream.close() - assert mock_client.capture.call_count == 1 - assert mock_client.capture.call_args.kwargs["properties"]["$ai_model"] == "gpt-4" + assert mock_client.capture_ai.call_count == 1 + assert mock_client.capture_ai.call_args.kwargs["properties"]["$ai_model"] == "gpt-4" diff --git a/posthog/test/ai/openai/test_openai.py b/posthog/test/ai/openai/test_openai.py index 0d22f393..01f6f842 100644 --- a/posthog/test/ai/openai/test_openai.py +++ b/posthog/test/ai/openai/test_openai.py @@ -393,9 +393,9 @@ def test_basic_completion(mock_client, mock_openai_response): ) assert response == mock_openai_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -440,9 +440,9 @@ def test_embeddings(mock_client, mock_embedding_response): ) assert response == mock_embedding_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -470,9 +470,9 @@ def test_groups(mock_client, mock_openai_response): ) assert response == mock_openai_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] assert call_args["groups"] == {"company": "test_company"} @@ -491,9 +491,9 @@ def test_privacy_mode_local(mock_client, mock_openai_response): ) assert response == mock_openai_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] is None assert props["$ai_output_choices"] is None @@ -514,9 +514,9 @@ def test_privacy_mode_global(mock_client, mock_openai_response): ) assert response == mock_openai_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] is None assert props["$ai_output_choices"] is None @@ -533,9 +533,9 @@ def test_error(mock_client, mock_openai_response): model="gpt-4", messages=[{"role": "user", "content": "Hello"}] ) - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_is_error"] is True assert props["$ai_error"] == "Test error" @@ -555,9 +555,9 @@ def test_cached_tokens(mock_client, mock_openai_response_with_cached_tokens): ) assert response == mock_openai_response_with_cached_tokens - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -594,9 +594,9 @@ def test_null_token_details_do_not_crash( ) assert response == mock_openai_response_with_null_token_details - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input_tokens"] == 20 @@ -630,9 +630,9 @@ def test_tool_calls(mock_client, mock_openai_response_with_tool_calls): ) assert response == mock_openai_response_with_tool_calls - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -701,9 +701,9 @@ def test_tool_calls_only_no_content(mock_client, mock_openai_response_tool_calls ) assert response == mock_openai_response_tool_calls_only - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -756,9 +756,9 @@ def test_responses_api_tool_calls(mock_client, mock_responses_api_with_tool_call ) assert response == mock_responses_api_with_tool_calls - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -825,9 +825,9 @@ def test_streaming_with_tool_calls(mock_client, streaming_tool_call_chunks): assert chunks == streaming_tool_call_chunks # Verify the capture was called with the right arguments - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -901,9 +901,9 @@ def test_responses_api(mock_client, mock_openai_response_with_responses_api): posthog_properties={"foo": "bar"}, ) assert response == mock_openai_response_with_responses_api - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -966,9 +966,9 @@ def test_responses_parse(mock_client, mock_parsed_response): ) assert response == mock_parsed_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -1018,9 +1018,9 @@ def test_chat_completions_parse(mock_client, mock_openai_response): assert response == mock_openai_response assert mock_parse.call_count == 1 assert "posthog_distinct_id" not in mock_parse.call_args.kwargs - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -1059,9 +1059,9 @@ async def test_async_chat_completions_parse(mock_client, mock_openai_response): assert response == mock_openai_response mock_parse.assert_awaited_once() assert "posthog_distinct_id" not in mock_parse.call_args.kwargs - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -1184,9 +1184,9 @@ def mock_streaming_response(**kwargs): assert "stream_options" not in captured_kwargs # Verify capture was called - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify tokens are captured correctly from response.usage (not 0) @@ -1249,8 +1249,8 @@ async def chunk_iterable(): assert len(chunks) == len(streaming_tool_call_chunks) assert chunks == streaming_tool_call_chunks - assert mock_client.capture.call_count == 1 - call_args = mock_client.capture.call_args[1] + assert mock_client.capture_ai.call_count == 1 + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -1317,8 +1317,8 @@ async def chunk_iterable(): kwargs = captured_kwargs["kwargs"] assert "stream_options" not in kwargs - assert mock_client.capture.call_count == 1 - call_args = mock_client.capture.call_args[1] + assert mock_client.capture_ai.call_count == 1 + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -1347,9 +1347,9 @@ async def test_async_embeddings_create(mock_client, mock_embedding_response): assert response == mock_embedding_response assert mock_create.await_count == 1 - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -1398,9 +1398,9 @@ def test_tool_definition(mock_client, mock_openai_response): ) assert response == mock_openai_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -1462,9 +1462,9 @@ def __init__(self): ) assert response == mock_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify web search count is detected (binary detection) @@ -1521,9 +1521,9 @@ def mock_create_call(**kwargs): ) assert result == mock_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify exact web search count @@ -1608,9 +1608,9 @@ def test_streaming_with_web_search(mock_client, streaming_web_search_chunks): # Verify the chunks were returned correctly assert len(chunks) == 3 - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify web search count is captured (binary detection = 1) @@ -1644,9 +1644,9 @@ def test_streaming_with_web_search_on_non_usage_chunk( # Verify the chunks were returned correctly assert len(chunks) == 3 - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify web search count is captured even though citations were on first chunk @@ -1700,9 +1700,9 @@ async def mock_create(self, **kwargs): ) assert response == mock_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify web search count is captured (binary detection = 1) @@ -1749,9 +1749,9 @@ async def chunk_iterable(): # Verify the chunks were returned correctly assert len(chunks) == 3 - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify web search count is captured (binary detection = 1) @@ -1816,8 +1816,8 @@ def test_streaming_chat_extracts_model_from_chunk_when_not_in_kwargs(mock_client # Consume the generator list(response_generator) - assert mock_client.capture.call_count == 1 - call_args = mock_client.capture.call_args[1] + assert mock_client.capture_ai.call_count == 1 + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Model should be extracted from chunk, not kwargs @@ -1861,7 +1861,7 @@ def test_streaming_chat_prefers_kwargs_model_over_chunk_model(mock_client): list(response_generator) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # kwargs model should take precedence @@ -1909,7 +1909,7 @@ def test_streaming_responses_api_extracts_model_from_response_object(mock_client list(response_generator) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Model should be extracted from chunk.response.model @@ -1983,7 +1983,7 @@ def test_streaming_responses_api_captures_usage_and_output_when_incomplete( ) list(response_generator) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_stop_reason"] == "max_output_tokens" assert props["$ai_input_tokens"] == 20 assert props["$ai_output_tokens"] == 16 @@ -2033,7 +2033,7 @@ def test_non_streaming_extracts_model_from_response(mock_client): ) assert response == mock_response - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Model should be extracted from response.model @@ -2093,7 +2093,7 @@ def test_non_streaming_responses_api_extracts_model_from_response(mock_client): ) assert response == mock_response - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Model should be extracted from response.model @@ -2139,7 +2139,7 @@ def test_non_streaming_returns_none_when_no_model(mock_client): posthog_distinct_id="test-id", ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Should be None, NOT "unknown" (to avoid incorrect cost matching) @@ -2178,7 +2178,7 @@ def test_streaming_falls_back_to_unknown_when_no_model(mock_client): list(response_generator) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Should fall back to "unknown" @@ -2231,7 +2231,7 @@ async def chunk_iterable(): async for _ in response_stream: pass - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_model"] == "gpt-4o-async-stored" @@ -2281,7 +2281,7 @@ async def chunk_iterable(): async for _ in response_stream: pass - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_model"] == "gpt-4o-mini-async-stored" @@ -2300,7 +2300,7 @@ def test_integration_stop_reason(mock_client): posthog_distinct_id="test-id", ) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_stop_reason"] in ("stop", "length") assert props["$ai_provider"] == "openai" assert props["$ai_input_tokens"] > 0 @@ -2334,8 +2334,8 @@ async def mock_create(self, **kwargs): chunks.append(chunk) assert chunks == streaming_tool_call_chunks - assert mock_client.capture.call_count == 1 - call_args = mock_client.capture.call_args[1] + assert mock_client.capture_ai.call_count == 1 + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["event"] == "$ai_generation" assert props["$ai_provider"] == "openai" @@ -2369,7 +2369,7 @@ async def mock_create(self, **kwargs): received = [c async for c in stream] assert received == [chunk] - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 @pytest.mark.asyncio @@ -2400,25 +2400,7 @@ async def mock_create(self, **kwargs): break assert source.closed is True - assert mock_client.capture.call_count == 1 - - -def test_ai_lane_client_routes_through_capture_ai(mock_client, mock_openai_response): - mock_client.enable_full_ai_capture = True - with patch( - "openai.resources.chat.completions.Completions.create", - return_value=mock_openai_response, - ): - client = OpenAI(api_key="test-key", posthog_client=mock_client) - client.chat.completions.create( - model="gpt-4", - messages=[{"role": "user", "content": "Hello"}], - posthog_distinct_id="test-id", - ) - - mock_client.capture.assert_not_called() assert mock_client.capture_ai.call_count == 1 - assert mock_client.capture_ai.call_args[1]["event"] == "$ai_generation" def test_multimodal_client_skips_media_redaction(mock_client, mock_openai_response): @@ -2481,8 +2463,8 @@ def test_provider_override_chat_completions(mock_client, mock_openai_response): # The override must never reach the underlying OpenAI request. assert "posthog_provider_override" not in mock_create.call_args.kwargs - assert mock_client.capture.call_count == 1 - call_args = mock_client.capture.call_args[1] + assert mock_client.capture_ai.call_count == 1 + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_provider"] == "deepseek" @@ -2507,8 +2489,8 @@ def test_provider_override_omitted_defaults_to_openai( posthog_distinct_id="test-id", ) - assert mock_client.capture.call_count == 1 - props = mock_client.capture.call_args[1]["properties"] + assert mock_client.capture_ai.call_count == 1 + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_provider"] == "openai" @@ -2549,8 +2531,8 @@ def test_provider_override_streaming_chat_completions(mock_client): list(response_generator) assert "posthog_provider_override" not in mock_create.call_args.kwargs - assert mock_client.capture.call_count == 1 - props = mock_client.capture.call_args[1]["properties"] + assert mock_client.capture_ai.call_count == 1 + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_provider"] == "groq" assert props["$ai_model"] == "gpt-4" @@ -2572,9 +2554,9 @@ def test_provider_override_responses_api( assert response == mock_openai_response_with_responses_api assert "posthog_provider_override" not in mock_create.call_args.kwargs - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_provider"] == "xai" assert props["$ai_model"] == "gpt-4o-mini" @@ -2595,9 +2577,9 @@ def test_provider_override_embeddings(mock_client, mock_embedding_response): assert response == mock_embedding_response assert "posthog_provider_override" not in mock_create.call_args.kwargs - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_provider"] == "mistral" assert props["$ai_model"] == "text-embedding-3-small" @@ -2620,9 +2602,9 @@ def test_provider_override_chat_completions_parse(mock_client, mock_openai_respo assert response == mock_openai_response assert mock_parse.call_count == 1 assert "posthog_provider_override" not in mock_parse.call_args.kwargs - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_provider"] == "cerebras" assert props["$ai_model"] == "gpt-4" @@ -2649,9 +2631,9 @@ async def test_async_provider_override_chat_completions( assert response == mock_openai_response mock_create.assert_awaited_once() assert "posthog_provider_override" not in mock_create.call_args.kwargs - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_provider"] == "together" assert props["$ai_model"] == "gpt-4" @@ -2675,8 +2657,8 @@ async def test_async_provider_override_omitted_defaults_to_openai( posthog_distinct_id="test-id", ) - assert mock_client.capture.call_count == 1 - props = mock_client.capture.call_args[1]["properties"] + assert mock_client.capture_ai.call_count == 1 + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_provider"] == "openai" @@ -2715,8 +2697,8 @@ async def chunk_iterable(): assert "posthog_provider_override" not in captured_kwargs["kwargs"] assert len(chunks) == len(streaming_tool_call_chunks) - assert mock_client.capture.call_count == 1 - props = mock_client.capture.call_args[1]["properties"] + assert mock_client.capture_ai.call_count == 1 + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_provider"] == "fireworks" assert props["$ai_model"] == "gpt-4" @@ -2738,9 +2720,9 @@ async def test_async_provider_override_embeddings(mock_client, mock_embedding_re assert response == mock_embedding_response assert mock_create.await_count == 1 assert "posthog_provider_override" not in mock_create.call_args.kwargs - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_provider"] == "perplexity" assert props["$ai_model"] == "text-embedding-3-small" @@ -2760,7 +2742,7 @@ def test_served_service_tier_lands_in_model_parameters( posthog_distinct_id="test-id", ) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_model_parameters"]["service_tier"] == "flex" assert props["$ai_service_tier"] == "flex" @@ -2777,7 +2759,7 @@ def test_response_without_service_tier_omits_it(mock_client, mock_openai_respons posthog_distinct_id="test-id", ) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] assert "service_tier" not in props["$ai_model_parameters"] assert "$ai_service_tier" not in props diff --git a/posthog/test/ai/openai_agents/test_processor.py b/posthog/test/ai/openai_agents/test_processor.py index 69745413..203434ab 100644 --- a/posthog/test/ai/openai_agents/test_processor.py +++ b/posthog/test/ai/openai_agents/test_processor.py @@ -109,7 +109,7 @@ 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() + mock_client.capture_ai.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): @@ -117,8 +117,8 @@ def test_on_trace_end_captures_ai_trace(self, processor, mock_client, mock_trace 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] + mock_client.capture_ai.assert_called_once() + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["event"] == "$ai_trace" assert call_kwargs["distinct_id"] == "test-user" @@ -137,10 +137,10 @@ def test_group_id_becomes_session_id( mock_span.span_data = GenerationSpanData(model="gpt-4o") processor.on_span_start(mock_span) processor.on_span_end(mock_span) - span_kwargs = mock_client.capture.call_args[1] + span_kwargs = mock_client.capture_ai.call_args[1] processor.on_trace_end(mock_trace) - trace_kwargs = mock_client.capture.call_args[1] + trace_kwargs = mock_client.capture_ai.call_args[1] assert span_kwargs["properties"]["$ai_session_id"] == "group_123" assert span_kwargs["properties"]["$ai_group_id"] == "group_123" @@ -156,7 +156,7 @@ def test_personless_mode_when_no_distinct_id(self, mock_client, mock_trace): processor.on_trace_start(mock_trace) processor.on_trace_end(mock_trace) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["options"] == {"process_person_profile": False} # Should fallback to trace_id as the distinct_id assert call_kwargs["distinct_id"] == mock_trace.trace_id @@ -170,7 +170,7 @@ def test_personless_mode_for_spans_when_no_distinct_id( ) processor.on_trace_start(mock_trace) - mock_client.capture.reset_mock() + mock_client.capture_ai.reset_mock() span_data = GenerationSpanData(model="gpt-4o") mock_span.span_data = span_data @@ -178,7 +178,7 @@ def test_personless_mode_for_spans_when_no_distinct_id( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["options"] == {"process_person_profile": False} assert call_kwargs["distinct_id"] == mock_span.trace_id @@ -196,7 +196,7 @@ def resolver(trace): ) processor.on_trace_start(mock_trace) - mock_client.capture.reset_mock() + mock_client.capture_ai.reset_mock() span_data = GenerationSpanData(model="gpt-4o") mock_span.span_data = span_data @@ -204,7 +204,7 @@ def resolver(trace): processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["options"] == {"process_person_profile": False} assert call_kwargs["distinct_id"] == mock_span.trace_id @@ -218,7 +218,7 @@ def test_person_profile_when_distinct_id_provided(self, mock_client, mock_trace) processor.on_trace_start(mock_trace) processor.on_trace_end(mock_trace) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert "options" not in call_kwargs def test_on_trace_end_clears_metadata(self, processor, mock_client, mock_trace): @@ -229,7 +229,7 @@ def test_on_trace_end_clears_metadata(self, processor, mock_client, mock_trace): 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() + mock_client.capture_ai.assert_called_once() def test_on_span_start_tracks_time(self, processor, mock_span): """Test that on_span_start records start time.""" @@ -250,8 +250,8 @@ def test_generation_span_mapping(self, processor, mock_client, mock_span): 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] + mock_client.capture_ai.assert_called_once() + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["event"] == "$ai_generation" assert call_kwargs["properties"]["$ai_trace_id"] == "trace_123456789" @@ -285,7 +285,7 @@ def test_generation_span_with_reasoning_tokens( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["properties"]["$ai_reasoning_tokens"] == 400 def test_function_span_mapping(self, processor, mock_client, mock_span): @@ -300,7 +300,7 @@ def test_function_span_mapping(self, processor, mock_client, mock_span): processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["event"] == "$ai_span" assert call_kwargs["properties"]["$ai_span_name"] == "get_weather" @@ -323,7 +323,7 @@ def test_agent_span_mapping(self, processor, mock_client, mock_span): processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["event"] == "$ai_span" assert call_kwargs["properties"]["$ai_span_name"] == "CustomerServiceAgent" @@ -345,7 +345,7 @@ def test_handoff_span_mapping(self, processor, mock_client, mock_span): processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["event"] == "$ai_span" assert call_kwargs["properties"]["$ai_span_type"] == "handoff" @@ -367,7 +367,7 @@ def test_guardrail_span_mapping(self, processor, mock_client, mock_span): processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["event"] == "$ai_span" assert call_kwargs["properties"]["$ai_span_name"] == "ContentFilter" @@ -385,7 +385,7 @@ def test_custom_span_mapping(self, processor, mock_client, mock_span): processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["event"] == "$ai_span" assert call_kwargs["properties"]["$ai_span_name"] == "database_query" @@ -414,7 +414,7 @@ def test_privacy_mode_redacts_content(self, mock_client, mock_span): processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] # Content should be redacted assert call_kwargs["properties"]["$ai_input"] is None @@ -435,7 +435,7 @@ def test_client_privacy_mode_redacts_content(self, mock_client, mock_span): processor.on_span_start(mock_span) processor.on_span_end(mock_span) - properties = mock_client.capture.call_args.kwargs["properties"] + properties = mock_client.capture_ai.call_args.kwargs["properties"] assert properties["$ai_input"] is None assert properties["$ai_output_choices"] is None @@ -480,7 +480,7 @@ def test_generation_span_image_input_is_redacted( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] captured_input = call_kwargs["properties"]["$ai_input"] assert captured_input[0]["content"][0]["image_url"] == "[base64 image redacted]" @@ -534,7 +534,7 @@ def test_response_span_image_output_is_redacted( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] captured_output = call_kwargs["properties"]["$ai_output_choices"] assert ( captured_output[0]["content"][0]["image_url"] == "[base64 image redacted]" @@ -549,7 +549,7 @@ def test_error_handling_in_span(self, processor, mock_client, mock_span): processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["properties"]["$ai_is_error"] is True assert call_kwargs["properties"]["$ai_error"] == "Rate limit exceeded" @@ -567,7 +567,7 @@ def test_generation_span_includes_total_tokens( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["properties"]["$ai_total_tokens"] == 150 def test_error_type_categorization_model_behavior( @@ -584,7 +584,7 @@ def test_error_type_categorization_model_behavior( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["properties"]["$ai_error_type"] == "model_behavior_error" def test_error_type_categorization_user_error( @@ -598,7 +598,7 @@ def test_error_type_categorization_user_error( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["properties"]["$ai_error_type"] == "user_error" def test_error_type_categorization_input_guardrail( @@ -614,7 +614,7 @@ def test_error_type_categorization_input_guardrail( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert ( call_kwargs["properties"]["$ai_error_type"] == "input_guardrail_triggered" ) @@ -632,7 +632,7 @@ def test_error_type_categorization_output_guardrail( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert ( call_kwargs["properties"]["$ai_error_type"] == "output_guardrail_triggered" ) @@ -648,7 +648,7 @@ def test_error_type_categorization_max_turns( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.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): @@ -660,7 +660,7 @@ def test_error_type_categorization_unknown(self, processor, mock_client, mock_sp processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["properties"]["$ai_error_type"] == "unknown" def test_response_span_with_output_and_total_tokens( @@ -686,7 +686,7 @@ def test_response_span_with_output_and_total_tokens( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["event"] == "$ai_generation" assert call_kwargs["properties"]["$ai_total_tokens"] == 35 @@ -715,7 +715,7 @@ def test_response_span_includes_total_cost( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["properties"]["$ai_total_cost_usd"] == total_cost_usd def test_speech_span_with_pass_through_properties( @@ -735,7 +735,7 @@ def test_speech_span_with_pass_through_properties( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["event"] == "$ai_span" assert call_kwargs["properties"]["$ai_span_type"] == "speech" @@ -768,7 +768,7 @@ def test_transcription_span_with_pass_through_properties( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["event"] == "$ai_span" assert call_kwargs["properties"]["$ai_span_type"] == "transcription" @@ -799,7 +799,7 @@ def test_transcription_span_audio_input_redacted( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["properties"]["$ai_input"] == "[base64 audio redacted]" def test_transcription_span_audio_input_passthrough( @@ -834,7 +834,7 @@ def test_latency_calculation(self, processor, mock_client, 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] + call_kwargs = mock_client.capture_ai.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): @@ -848,7 +848,7 @@ def test_groups_included_in_events(self, mock_client, mock_trace, mock_span): processor.on_trace_start(mock_trace) processor.on_trace_end(mock_trace) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["groups"] == {"company": "acme", "team": "engineering"} def test_additional_properties_included(self, mock_client, mock_trace): @@ -862,7 +862,7 @@ def test_additional_properties_included(self, mock_client, mock_trace): processor.on_trace_start(mock_trace) processor.on_trace_end(mock_trace) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["properties"]["environment"] == "production" assert call_kwargs["properties"]["version"] == "1.0" @@ -889,7 +889,7 @@ def test_generation_span_with_no_usage(self, processor, mock_client, mock_span): processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert "$ai_input_tokens" not in call_kwargs["properties"] assert "$ai_output_tokens" not in call_kwargs["properties"] assert "$ai_total_tokens" not in call_kwargs["properties"] @@ -907,7 +907,7 @@ def test_generation_span_with_partial_usage( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["properties"]["$ai_input_tokens"] == 42 # The output side was never reported, so it is omitted rather than 0; # the total is the sum of the reported sides. @@ -928,7 +928,7 @@ def test_error_type_categorization_by_type_field_only( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["properties"]["$ai_error_type"] == "model_behavior_error" def test_distinct_id_resolved_from_trace_for_spans( @@ -946,7 +946,7 @@ def resolver(trace): # Start trace - this resolves and stores distinct_id processor.on_trace_start(mock_trace) - mock_client.capture.reset_mock() + mock_client.capture_ai.reset_mock() # End a span - should use the stored distinct_id from trace span_data = GenerationSpanData(model="gpt-4o") @@ -955,7 +955,7 @@ def resolver(trace): processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["distinct_id"] == "user-Test Workflow" def test_eviction_of_stale_entries(self, mock_client): @@ -1044,24 +1044,13 @@ def test_processor_groups_are_forwarded(self, mock_client): processor._capture_event("$ai_trace", {}) - assert mock_client.capture.call_args.kwargs["groups"] == {"company": "acme"} + assert mock_client.capture_ai.call_args.kwargs["groups"] == {"company": "acme"} def test_capture_errors_are_logged_and_suppressed(self, mock_client, caplog): - mock_client.capture.side_effect = RuntimeError("capture failed") + mock_client.capture_ai.side_effect = RuntimeError("capture failed") processor = PostHogTracingProcessor(client=mock_client) with caplog.at_level(logging.DEBUG, logger="posthog"): processor._capture_event("$ai_trace", {}) assert "Failed to capture PostHog event: capture failed" in caplog.text - - -def test_ai_lane_client_routes_through_capture_ai(mock_client, mock_trace): - mock_client.enable_full_ai_capture = True - processor = PostHogTracingProcessor(client=mock_client, distinct_id="test-user") - processor.on_trace_start(mock_trace) - processor.on_trace_end(mock_trace) - - mock_client.capture.assert_not_called() - mock_client.capture_ai.assert_called_once() - assert mock_client.capture_ai.call_args[1]["event"] == "$ai_trace" diff --git a/posthog/test/ai/test_token_reporting.py b/posthog/test/ai/test_token_reporting.py index 790a3a6d..e06c2f07 100644 --- a/posthog/test/ai/test_token_reporting.py +++ b/posthog/test/ai/test_token_reporting.py @@ -52,7 +52,7 @@ def test_token_counts_trace_back_to_a_provider_report( ): capture_streaming_event(mock_client, _event_data(usage_stats)) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] for key in ("$ai_input_tokens", "$ai_output_tokens"): if key in expected: assert props[key] == expected[key] @@ -85,7 +85,7 @@ def test_aux_token_fields_trace_back_to_a_provider_report( ): capture_streaming_event(mock_client, _event_data(usage_stats, provider=provider)) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] for key in ( "$ai_cache_read_input_tokens", "$ai_cache_creation_input_tokens", diff --git a/posthog/test/test_ai_capture_lane.py b/posthog/test/test_ai_capture_lane.py index 2234de1b..9034f874 100644 --- a/posthog/test/test_ai_capture_lane.py +++ b/posthog/test/test_ai_capture_lane.py @@ -83,14 +83,6 @@ def test_sync_mode_capture_ai_posts_single_event_batch_to_ai_endpoint(self): batch = sent_batch(mock_post) self.assertEqual([e["event"] for e in batch], ["$ai_generation"]) - def test_multimodal_client_routes_wrapper_captures_to_ai_lane(self): - client = self._client(enable_full_ai_capture=True) - with patch_capture_send("consumer") as mock_post: - _capture_ai_event(client, "$ai_generation", distinct_id="d") - client.flush() - by_path = _events_by_path(mock_post) - self.assertEqual(set(by_path), {_CAPTURE_AI_V1_PATH}) - def test_disabled_client_never_starts_ai_lane(self): client = Client(TEST_API_KEY, disabled=True) client.capture_ai("$ai_generation", distinct_id="d") @@ -464,10 +456,13 @@ def test_fork_rebuild_replaces_sync_mode_queues(self): class TestCaptureAiEventHelper(unittest.TestCase): - """`_capture_ai_event` rides the AI lane only when the client opted in.""" + """`_capture_ai_event` always rides the AI lane; `enable_full_ai_capture` only controls content.""" - def test_opted_in_routes_through_ai_lane(self): - client = Client(TEST_API_KEY, flush_interval=0.05, enable_full_ai_capture=True) + @parameterized.expand([("default", False), ("full_capture", True)]) + def test_routes_only_through_ai_lane(self, _name, full_capture): + client = Client( + TEST_API_KEY, flush_interval=0.05, enable_full_ai_capture=full_capture + ) with patch_capture_send("consumer") as mock_post: _capture_ai_event( client, @@ -483,19 +478,6 @@ def test_opted_in_routes_through_ai_lane(self): ) client.join() - def test_default_keeps_capture_path(self): - client = Client(TEST_API_KEY, flush_interval=0.05) - with patch_capture_send("consumer") as mock_post: - _capture_ai_event(client, "$ai_generation", distinct_id="d") - client.flush() - - self.assertEqual( - [call.kwargs["path"] for call in mock_post.call_args_list], - [_CAPTURE_V1_PATH], - ) - self.assertEqual(client._ai_lane.consumers, []) - client.join() - def test_adds_ai_library_identity_and_preserves_provider_and_model(self): client = mock.Mock() _capture_ai_event( @@ -505,59 +487,16 @@ def test_adds_ai_library_identity_and_preserves_provider_and_model(self): properties={"$ai_provider": "openai", "$ai_model": "gpt-4o"}, ) - properties = client.capture.call_args.kwargs["properties"] + properties = client.capture_ai.call_args.kwargs["properties"] self.assertEqual(properties["$ai_lib"], "posthog-ai") self.assertEqual(properties["$ai_lib_version"], VERSION) self.assertEqual(properties["$ai_provider"], "openai") self.assertEqual(properties["$ai_model"], "gpt-4o") - def test_default_mock_clients_keep_seeing_capture(self): - # Downstream test suites pass Mock clients into the wrappers; without - # the opt-in they must keep seeing plain `capture()` calls. - client = mock.Mock() - _capture_ai_event(client, "$ai_generation", distinct_id="d") - client.capture.assert_called_once_with( - event="$ai_generation", - distinct_id="d", - properties={"$ai_lib": "posthog-ai", "$ai_lib_version": VERSION}, - ) - client.capture_ai.assert_not_called() - - def test_opted_in_prefers_capture_ai(self): - client = mock.Mock(spec=["capture", "capture_ai", "enable_full_ai_capture"]) - client.enable_full_ai_capture = True - _capture_ai_event(client, "$ai_generation", distinct_id="d") - client.capture_ai.assert_called_once_with( - event="$ai_generation", - distinct_id="d", - properties={"$ai_lib": "posthog-ai", "$ai_lib_version": VERSION}, - ) - client.capture.assert_not_called() - - def test_opted_in_duck_typed_client_without_method_falls_back(self): + @parameterized.expand([("default", False), ("full_capture", True)]) + def test_client_without_capture_ai_falls_back_to_capture(self, _name, full_capture): client = mock.Mock(spec=["capture", "enable_full_ai_capture"]) - client.enable_full_ai_capture = True - _capture_ai_event(client, "$ai_generation", distinct_id="d") - client.capture.assert_called_once_with( - event="$ai_generation", - distinct_id="d", - properties={"$ai_lib": "posthog-ai", "$ai_lib_version": VERSION}, - ) - - def test_client_multimodal_flag_prefers_capture_ai(self): - client = mock.Mock(spec=["capture", "capture_ai", "enable_full_ai_capture"]) - client.enable_full_ai_capture = True - _capture_ai_event(client, "$ai_generation", distinct_id="d") - client.capture_ai.assert_called_once_with( - event="$ai_generation", - distinct_id="d", - properties={"$ai_lib": "posthog-ai", "$ai_lib_version": VERSION}, - ) - client.capture.assert_not_called() - - def test_client_multimodal_flag_off_keeps_capture(self): - client = mock.Mock(spec=["capture", "capture_ai", "enable_full_ai_capture"]) - client.enable_full_ai_capture = False + client.enable_full_ai_capture = full_capture _capture_ai_event(client, "$ai_generation", distinct_id="d") client.capture.assert_called_once_with( event="$ai_generation", @@ -571,7 +510,7 @@ class TestLanesRefuseWorkAfterShutdown(unittest.TestCase): afterwards, even a lazy AI lane that never started before shutdown.""" def test_late_ai_capture_after_shutdown_starts_nothing_and_sends_nothing(self): - client = Client(TEST_API_KEY, enable_full_ai_capture=True, flush_interval=0.05) + client = Client(TEST_API_KEY, flush_interval=0.05) client.shutdown() with patch_capture_send("consumer") as mock_post: _capture_ai_event(client, "$ai_generation", distinct_id="d")