From d127621439bc643ddcfc3484b302102cfc642b6f Mon Sep 17 00:00:00 2001 From: Bernat Torres Date: Mon, 31 Aug 2026 18:02:58 +0200 Subject: [PATCH 1/2] fix(ai): omit token counts the provider never reported Every integration wrote $ai_input_tokens: 0 and $ai_output_tokens: 0 when the provider never reported usage, so an interrupted stream priced as a free call downstream. Token counts are now written only when they trace back to a provider report: the shared capture paths, both embedding events, the LangChain callback and the Agents and Claude Agent SDK processors omit unreported counts, and the stream accumulators and response extractors no longer seed at zero. A zero the provider reported is still sent as 0, so zero keeps meaning a real report of nothing. LangChain cache and reasoning counts follow the same rule instead of coalescing to 0, and the Agents processor's $ai_total_tokens is the sum of the reported sides, omitted when neither side reported. Mirrors the posthog-js change (PostHog/posthog-js#4664). Ingestion distinguishes absent from zero since PostHog/posthog#90211 and older ingestion coalesces absent to 0 at read time, so this ships independently. Generated-By: PostHog Desktop Task-Id: ec0cc27f-fc40-4994-b153-2bfe74de1edf --- .../omit-unreported-token-counts.md | 5 ++ posthog/ai/anthropic/_anthropic_stream.py | 2 +- posthog/ai/anthropic/anthropic_converter.py | 17 +++--- posthog/ai/claude_agent_sdk/processor.py | 48 ++++++++++----- posthog/ai/gemini/_shared.py | 8 ++- posthog/ai/gemini/gemini.py | 2 +- posthog/ai/gemini/gemini_async.py | 2 +- posthog/ai/gemini/gemini_converter.py | 10 ++-- posthog/ai/langchain/callbacks.py | 25 +++++--- posthog/ai/openai/_embeddings.py | 6 +- posthog/ai/openai/openai_converter.py | 21 +++---- posthog/ai/openai_agents/processor.py | 51 +++++++++++----- posthog/ai/utils.py | 31 +++++++--- posthog/test/ai/gemini/test_gemini.py | 26 ++++++++- posthog/test/ai/gemini/test_gemini_async.py | 3 +- posthog/test/ai/langchain/test_callbacks.py | 21 ++++--- .../test/ai/openai_agents/test_processor.py | 12 ++-- posthog/test/ai/test_token_reporting.py | 58 +++++++++++++++++++ 18 files changed, 259 insertions(+), 89 deletions(-) create mode 100644 .sampo/changesets/omit-unreported-token-counts.md create mode 100644 posthog/test/ai/test_token_reporting.py diff --git a/.sampo/changesets/omit-unreported-token-counts.md b/.sampo/changesets/omit-unreported-token-counts.md new file mode 100644 index 000000000..5c4ae9af0 --- /dev/null +++ b/.sampo/changesets/omit-unreported-token-counts.md @@ -0,0 +1,5 @@ +--- +pypi/posthog: patch +--- + +Omit `$ai_input_tokens` and `$ai_output_tokens` when the provider never reported usage, instead of sending `0`, so an interrupted stream no longer looks like a free call. A zero reported by the provider is still sent, and zero keeps meaning a real report of nothing. Covers the OpenAI, Anthropic, Gemini, LangChain, OpenAI Agents and Claude Agent SDK integrations. diff --git a/posthog/ai/anthropic/_anthropic_stream.py b/posthog/ai/anthropic/_anthropic_stream.py index 3426e1d4b..4b94913ea 100644 --- a/posthog/ai/anthropic/_anthropic_stream.py +++ b/posthog/ai/anthropic/_anthropic_stream.py @@ -15,7 +15,7 @@ class _AnthropicStreamAccumulator: """Accumulates sync-neutral capture state from Anthropic stream events.""" def __init__(self) -> None: - self.usage_stats: TokenUsage = TokenUsage(input_tokens=0, output_tokens=0) + self.usage_stats: TokenUsage = TokenUsage() self.accumulated_content = "" self.content_blocks: List[StreamingContentBlock] = [] self.tools_in_progress: Dict[str, ToolInProgress] = {} diff --git a/posthog/ai/anthropic/anthropic_converter.py b/posthog/ai/anthropic/anthropic_converter.py index ab5503fc0..650c5587f 100644 --- a/posthog/ai/anthropic/anthropic_converter.py +++ b/posthog/ai/anthropic/anthropic_converter.py @@ -252,13 +252,16 @@ def extract_anthropic_usage_from_response(response: Any) -> TokenUsage: Returns: TokenUsage with standardized usage """ - if not hasattr(response, "usage"): - return TokenUsage(input_tokens=0, output_tokens=0) - - result = TokenUsage( - input_tokens=getattr(response.usage, "input_tokens", 0), - output_tokens=getattr(response.usage, "output_tokens", 0), - ) + if getattr(response, "usage", None) is None: + return TokenUsage() + + result = TokenUsage() + input_tokens = getattr(response.usage, "input_tokens", None) + if input_tokens is not None: + result["input_tokens"] = input_tokens + output_tokens = getattr(response.usage, "output_tokens", None) + if output_tokens is not None: + result["output_tokens"] = output_tokens if hasattr(response.usage, "cache_read_input_tokens"): cache_read = response.usage.cache_read_input_tokens diff --git a/posthog/ai/claude_agent_sdk/processor.py b/posthog/ai/claude_agent_sdk/processor.py index cc53726e4..1d45b61c2 100644 --- a/posthog/ai/claude_agent_sdk/processor.py +++ b/posthog/ai/claude_agent_sdk/processor.py @@ -45,10 +45,12 @@ class _GenerationData: """Data accumulated for a single LLM generation (one API call).""" model: Optional[str] = None - input_tokens: int = 0 - output_tokens: int = 0 - cache_read_input_tokens: int = 0 - cache_creation_input_tokens: int = 0 + # None when the provider never reported a count: absent means unknown, + # 0 is a report of nothing. + input_tokens: Optional[int] = None + output_tokens: Optional[int] = None + cache_read_input_tokens: Optional[int] = None + cache_creation_input_tokens: Optional[int] = None raw_usage: Optional[Dict[str, Any]] = None start_time: float = 0.0 end_time: float = 0.0 @@ -79,13 +81,11 @@ def process_stream_event(self, event: "StreamEvent") -> None: message = raw.get("message", {}) self._current.model = message.get("model") usage = message.get("usage", {}) - self._current.input_tokens = usage.get("input_tokens", 0) - self._current.output_tokens = usage.get("output_tokens", 0) - self._current.cache_read_input_tokens = usage.get( - "cache_read_input_tokens", 0 - ) + self._current.input_tokens = usage.get("input_tokens") + self._current.output_tokens = usage.get("output_tokens") + self._current.cache_read_input_tokens = usage.get("cache_read_input_tokens") self._current.cache_creation_input_tokens = usage.get( - "cache_creation_input_tokens", 0 + "cache_creation_input_tokens" ) self._current.raw_usage = dict(usage) @@ -410,8 +410,16 @@ def _emit_generation( "$ai_provider": "anthropic", "$ai_framework": "claude-agent-sdk", "$ai_model": gen.model, - "$ai_input_tokens": gen.input_tokens, - "$ai_output_tokens": gen.output_tokens, + **( + {"$ai_input_tokens": gen.input_tokens} + if gen.input_tokens is not None + else {} + ), + **( + {"$ai_output_tokens": gen.output_tokens} + if gen.output_tokens is not None + else {} + ), "$ai_latency": latency, **extra_props, } @@ -472,8 +480,16 @@ def _emit_generation_from_result( "$ai_provider": "anthropic", "$ai_framework": "claude-agent-sdk", "$ai_model": model, - "$ai_input_tokens": usage.get("input_tokens", 0), - "$ai_output_tokens": usage.get("output_tokens", 0), + **( + {"$ai_input_tokens": usage["input_tokens"]} + if usage.get("input_tokens") is not None + else {} + ), + **( + {"$ai_output_tokens": usage["output_tokens"]} + if usage.get("output_tokens") is not None + else {} + ), "$ai_latency": result.duration_api_ms / 1000.0 if result.duration_api_ms else 0, @@ -494,8 +510,8 @@ def _emit_generation_from_result( finalize_ai_content(output_choices, self._client), ) - cache_read = usage.get("cache_read_input_tokens", 0) - cache_creation = usage.get("cache_creation_input_tokens", 0) + cache_read = usage.get("cache_read_input_tokens") + cache_creation = usage.get("cache_creation_input_tokens") if cache_read: properties["$ai_cache_read_input_tokens"] = cache_read if cache_creation: diff --git a/posthog/ai/gemini/_shared.py b/posthog/ai/gemini/_shared.py index 9ac080c8b..95ebf019d 100644 --- a/posthog/ai/gemini/_shared.py +++ b/posthog/ai/gemini/_shared.py @@ -195,7 +195,9 @@ def _capture_embedding_outcome( error: Optional[Exception], latency: float, ) -> None: - input_tokens = extract_gemini_embedding_token_count(response) if response else 0 + input_tokens = ( + extract_gemini_embedding_token_count(response) if response else None + ) event_properties = { "$ai_provider": "gemini", "$ai_model": model, @@ -207,7 +209,9 @@ def _capture_embedding_outcome( "$ai_http_status": ( getattr(error, "status_code", 0) if error is not None else 200 ), - "$ai_input_tokens": input_tokens, + # Omitted when the provider never reported a count: absent means + # unknown, 0 is a report of nothing. + **({"$ai_input_tokens": input_tokens} if input_tokens is not None else {}), "$ai_latency": latency, "$ai_trace_id": trace_id, "$ai_base_url": self._base_url, diff --git a/posthog/ai/gemini/gemini.py b/posthog/ai/gemini/gemini.py index d8e977e9f..da76cb3db 100644 --- a/posthog/ai/gemini/gemini.py +++ b/posthog/ai/gemini/gemini.py @@ -216,7 +216,7 @@ def _generate_content_streaming( **kwargs: Any, ): start_time = time.time() - usage_stats: TokenUsage = TokenUsage(input_tokens=0, output_tokens=0) + usage_stats: TokenUsage = TokenUsage() accumulated_content = [] stop_reason: Optional[str] = None diff --git a/posthog/ai/gemini/gemini_async.py b/posthog/ai/gemini/gemini_async.py index 78ee7dbee..bf05a68d8 100644 --- a/posthog/ai/gemini/gemini_async.py +++ b/posthog/ai/gemini/gemini_async.py @@ -217,7 +217,7 @@ async def _generate_content_streaming( **kwargs: Any, ): start_time = time.time() - usage_stats: TokenUsage = TokenUsage(input_tokens=0, output_tokens=0) + usage_stats: TokenUsage = TokenUsage() accumulated_content = [] stop_reason: Optional[str] = None diff --git a/posthog/ai/gemini/gemini_converter.py b/posthog/ai/gemini/gemini_converter.py index b5548be42..528c3db9f 100644 --- a/posthog/ai/gemini/gemini_converter.py +++ b/posthog/ai/gemini/gemini_converter.py @@ -546,7 +546,7 @@ def extract_gemini_usage_from_response(response: Any) -> TokenUsage: TokenUsage with standardized usage statistics """ if not hasattr(response, "usage_metadata") or not response.usage_metadata: - return TokenUsage(input_tokens=0, output_tokens=0) + return TokenUsage() usage = _extract_usage_from_metadata(response.usage_metadata) @@ -715,17 +715,19 @@ def format_gemini_streaming_output( return [{"role": "assistant", "content": [{"type": "text", "text": ""}]}] -def extract_gemini_embedding_token_count(response) -> int: +def extract_gemini_embedding_token_count(response) -> Optional[int]: """ Extract total token count from a Gemini embed_content response. Token counts are only available per-embedding via Vertex AI's statistics.token_count. - Returns 0 if no token counts are available. + Returns None when no embedding carried a token count. """ total = 0 + reported = False if hasattr(response, "embeddings") and response.embeddings: for embedding in response.embeddings: if hasattr(embedding, "statistics") and embedding.statistics: token_count = getattr(embedding.statistics, "token_count", None) if token_count is not None: total += int(token_count) - return total + reported = True + return total if reported else None diff --git a/posthog/ai/langchain/callbacks.py b/posthog/ai/langchain/callbacks.py index c659c62af..91fdece52 100644 --- a/posthog/ai/langchain/callbacks.py +++ b/posthog/ai/langchain/callbacks.py @@ -664,11 +664,16 @@ def _capture_generation( else: # Add usage usage = _parse_usage(output, run.provider, run.model) - event_properties["$ai_input_tokens"] = usage.input_tokens - event_properties["$ai_output_tokens"] = usage.output_tokens - event_properties["$ai_cache_creation_input_tokens"] = ( - usage.cache_write_tokens - ) + # Omitted when the provider never reported a count: absent means + # unknown, 0 is a report of nothing. + if usage.input_tokens is not None: + event_properties["$ai_input_tokens"] = usage.input_tokens + if usage.output_tokens is not None: + event_properties["$ai_output_tokens"] = usage.output_tokens + if usage.cache_write_tokens is not None: + event_properties["$ai_cache_creation_input_tokens"] = ( + usage.cache_write_tokens + ) if ( usage.cache_write_5m_tokens is not None and usage.cache_write_1h_tokens is not None @@ -679,8 +684,12 @@ def _capture_generation( event_properties["$ai_cache_creation_1h_input_tokens"] = ( usage.cache_write_1h_tokens ) - event_properties["$ai_cache_read_input_tokens"] = usage.cache_read_tokens - event_properties["$ai_reasoning_tokens"] = usage.reasoning_tokens + if usage.cache_read_tokens is not None: + event_properties["$ai_cache_read_input_tokens"] = ( + usage.cache_read_tokens + ) + if usage.reasoning_tokens is not None: + event_properties["$ai_reasoning_tokens"] = usage.reasoning_tokens # Generation results generation_result = output.generations[-1] @@ -875,7 +884,7 @@ def _parse_usage_model( } normalized_usage = ModelUsage( **{ - dataclass_key: parsed_usage.get(mapped_key) or 0 + dataclass_key: parsed_usage.get(mapped_key) for mapped_key, dataclass_key in field_mapping.items() }, cache_write_5m_tokens=parsed_usage.get("cache_write_5m"), diff --git a/posthog/ai/openai/_embeddings.py b/posthog/ai/openai/_embeddings.py index 834e8bc57..8cdb04d33 100644 --- a/posthog/ai/openai/_embeddings.py +++ b/posthog/ai/openai/_embeddings.py @@ -18,7 +18,7 @@ def _capture_embedding_event( ) -> None: """Build and capture telemetry shared by sync and async embedding wrappers.""" usage = getattr(response, "usage", None) - input_tokens = getattr(usage, "prompt_tokens", 0) if usage else 0 + input_tokens = getattr(usage, "prompt_tokens", None) if usage else None event_properties = { "$ai_provider": "openai", @@ -29,7 +29,9 @@ def _capture_embedding_event( finalize_ai_content(request_kwargs.get("input"), posthog_client), ), "$ai_http_status": 200, - "$ai_input_tokens": input_tokens, + # Omitted when the provider never reported a count: absent means + # unknown, 0 is a report of nothing. + **({"$ai_input_tokens": input_tokens} if input_tokens is not None else {}), "$ai_latency": latency, "$ai_trace_id": trace_id, "$ai_base_url": str(base_url), diff --git a/posthog/ai/openai/openai_converter.py b/posthog/ai/openai/openai_converter.py index 1ce209003..ecd8c8c90 100644 --- a/posthog/ai/openai/openai_converter.py +++ b/posthog/ai/openai/openai_converter.py @@ -460,13 +460,13 @@ def extract_openai_usage_from_response(response: Any) -> TokenUsage: Returns: TokenUsage with standardized usage statistics """ - if not hasattr(response, "usage"): - return TokenUsage(input_tokens=0, output_tokens=0) + if not hasattr(response, "usage") or not response.usage: + return TokenUsage() - cached_tokens = 0 - input_tokens = 0 - output_tokens = 0 - reasoning_tokens = 0 + cached_tokens = None + input_tokens = None + output_tokens = None + reasoning_tokens = None # Responses API format if hasattr(response.usage, "input_tokens"): @@ -496,10 +496,11 @@ def extract_openai_usage_from_response(response: Any) -> TokenUsage: ): reasoning_tokens = response.usage.completion_tokens_details.reasoning_tokens - result = TokenUsage( - input_tokens=input_tokens, - output_tokens=output_tokens, - ) + result = TokenUsage() + if input_tokens is not None: + result["input_tokens"] = input_tokens + if output_tokens is not None: + result["output_tokens"] = output_tokens if cached_tokens is not None and cached_tokens > 0: result["cache_read_input_tokens"] = cached_tokens diff --git a/posthog/ai/openai_agents/processor.py b/posthog/ai/openai_agents/processor.py index d5418089f..8bdea6c7e 100644 --- a/posthog/ai/openai_agents/processor.py +++ b/posthog/ai/openai_agents/processor.py @@ -478,10 +478,14 @@ def _handle_generation_span( """Handle LLM generation spans - maps to $ai_generation event.""" # Extract token usage usage = span_data.usage or {} - input_tokens = usage.get("input_tokens") or usage.get("prompt_tokens") or 0 - output_tokens = ( - usage.get("output_tokens") or usage.get("completion_tokens") or 0 - ) + # None when the span never reported a count: absent means unknown, + # 0 is a report of nothing. + input_tokens = usage.get("input_tokens") + if input_tokens is None: + input_tokens = usage.get("prompt_tokens") + output_tokens = usage.get("output_tokens") + if output_tokens is None: + output_tokens = usage.get("completion_tokens") # Extract model config parameters model_config = span_data.model_config or {} @@ -510,9 +514,18 @@ def _handle_generation_span( _ensure_serializable(span_data.output), self._client ) ), - "$ai_input_tokens": input_tokens, - "$ai_output_tokens": output_tokens, - "$ai_total_tokens": (input_tokens or 0) + (output_tokens or 0), + **({"$ai_input_tokens": input_tokens} if input_tokens is not None else {}), + **( + {"$ai_output_tokens": output_tokens} + if output_tokens is not None + else {} + ), + # Sum of the reported sides; omitted when neither side was reported. + **( + {"$ai_total_tokens": (input_tokens or 0) + (output_tokens or 0)} + if input_tokens is not None or output_tokens is not None + else {} + ), } # Add optional token fields if present @@ -661,11 +674,10 @@ def _handle_response_span( # Try to extract usage from response usage = getattr(response, "usage", None) if response else None total_cost_usd = getattr(usage, "cost", None) if usage else None - input_tokens = 0 - output_tokens = 0 - if usage: - input_tokens = getattr(usage, "input_tokens", 0) or 0 - output_tokens = getattr(usage, "output_tokens", 0) or 0 + # None when the response never reported a count: absent means unknown, + # 0 is a report of nothing. + input_tokens = getattr(usage, "input_tokens", None) if usage else None + output_tokens = getattr(usage, "output_tokens", None) if usage else None # Try to extract model from response model = getattr(response, "model", None) if response else None @@ -679,9 +691,18 @@ def _handle_response_span( "$ai_input": self._with_privacy_mode( finalize_ai_content(_ensure_serializable(span_data.input), self._client) ), - "$ai_input_tokens": input_tokens, - "$ai_output_tokens": output_tokens, - "$ai_total_tokens": input_tokens + output_tokens, + **({"$ai_input_tokens": input_tokens} if input_tokens is not None else {}), + **( + {"$ai_output_tokens": output_tokens} + if output_tokens is not None + else {} + ), + # Sum of the reported sides; omitted when neither side was reported. + **( + {"$ai_total_tokens": (input_tokens or 0) + (output_tokens or 0)} + if input_tokens is not None or output_tokens is not None + else {} + ), } if total_cost_usd is not None: diff --git a/posthog/ai/utils.py b/posthog/ai/utils.py index 776539aa1..c456c0b72 100644 --- a/posthog/ai/utils.py +++ b/posthog/ai/utils.py @@ -281,7 +281,7 @@ def get_usage(response, provider: str) -> TokenUsage: return extract_gemini_usage_from_response(response) - return TokenUsage(input_tokens=0, output_tokens=0) + return TokenUsage() def format_response(response, provider: str): @@ -490,8 +490,14 @@ def call_llm_and_track_usage( ), ) tag("$ai_http_status", http_status) - tag("$ai_input_tokens", usage.get("input_tokens", 0)) - tag("$ai_output_tokens", usage.get("output_tokens", 0)) + # Omitted when the provider never reported a count: absent means + # unknown, 0 is a report of nothing. + input_tokens = usage.get("input_tokens") + if input_tokens is not None: + tag("$ai_input_tokens", input_tokens) + output_tokens = usage.get("output_tokens") + if output_tokens is not None: + tag("$ai_output_tokens", output_tokens) tag("$ai_latency", latency) tag("$ai_trace_id", posthog_trace_id) tag("$ai_base_url", str(base_url)) @@ -647,8 +653,14 @@ async def call_llm_and_track_usage_async( ), ) tag("$ai_http_status", http_status) - tag("$ai_input_tokens", usage.get("input_tokens", 0)) - tag("$ai_output_tokens", usage.get("output_tokens", 0)) + # Omitted when the provider never reported a count: absent means + # unknown, 0 is a report of nothing. + input_tokens = usage.get("input_tokens") + if input_tokens is not None: + tag("$ai_input_tokens", input_tokens) + output_tokens = usage.get("output_tokens") + if output_tokens is not None: + tag("$ai_output_tokens", output_tokens) tag("$ai_latency", latency) tag("$ai_trace_id", posthog_trace_id) tag("$ai_base_url", str(base_url)) @@ -765,6 +777,11 @@ def capture_streaming_event( """ trace_id = event_data.get("trace_id") or str(uuid.uuid4()) + # Omitted when the provider never reported a count: absent means unknown, + # 0 is a report of nothing. An interrupted stream often reports neither. + input_tokens = event_data["usage_stats"].get("input_tokens") + output_tokens = event_data["usage_stats"].get("output_tokens") + # Build base event properties event_properties = { "$ai_provider": event_data["provider"], @@ -781,8 +798,8 @@ def capture_streaming_event( finalize_ai_content(event_data["formatted_output"], ph_client), ), "$ai_http_status": 200, - "$ai_input_tokens": event_data["usage_stats"].get("input_tokens", 0), - "$ai_output_tokens": event_data["usage_stats"].get("output_tokens", 0), + **({"$ai_input_tokens": input_tokens} if input_tokens is not None else {}), + **({"$ai_output_tokens": output_tokens} if output_tokens is not None else {}), "$ai_latency": event_data["latency"], "$ai_trace_id": trace_id, "$ai_base_url": str(event_data["base_url"]), diff --git a/posthog/test/ai/gemini/test_gemini.py b/posthog/test/ai/gemini/test_gemini.py index f938f618f..f2a801e68 100644 --- a/posthog/test/ai/gemini/test_gemini.py +++ b/posthog/test/ai/gemini/test_gemini.py @@ -1312,7 +1312,31 @@ def test_embed_content_without_token_counts( ) props = mock_client.capture.call_args[1]["properties"] - assert props["$ai_input_tokens"] == 0 + # No embedding carried a token count, so the property is omitted, not 0. + assert "$ai_input_tokens" not in props + + +def test_streaming_without_usage_omits_token_counts( + mock_client, mock_google_genai_client +): + """A stream that never reported usage omits the counts instead of sending 0.""" + chunk = MagicMock() + chunk.text = "Hello" + chunk.usage_metadata = None + + mock_google_genai_client.models.generate_content_stream.return_value = iter([chunk]) + + client = Client(api_key="test-key", posthog_client=mock_client) + response = client.models.generate_content_stream( + model="gemini-2.0-flash", + contents=["Write a short story"], + posthog_distinct_id="test-id", + ) + list(response) + + props = mock_client.capture.call_args[1]["properties"] + assert "$ai_input_tokens" not in props + assert "$ai_output_tokens" not in props def test_embed_content_privacy_mode( diff --git a/posthog/test/ai/gemini/test_gemini_async.py b/posthog/test/ai/gemini/test_gemini_async.py index 53f85b264..8dc51ca34 100644 --- a/posthog/test/ai/gemini/test_gemini_async.py +++ b/posthog/test/ai/gemini/test_gemini_async.py @@ -954,7 +954,8 @@ async def test_async_embed_content_without_token_counts( ) props = mock_client.capture.call_args[1]["properties"] - assert props["$ai_input_tokens"] == 0 + # No embedding carried a token count, so the property is omitted, not 0. + assert "$ai_input_tokens" not in props async def test_async_embed_content_privacy_mode( diff --git a/posthog/test/ai/langchain/test_callbacks.py b/posthog/test/ai/langchain/test_callbacks.py index 577178fcd..11f021099 100644 --- a/posthog/test/ai/langchain/test_callbacks.py +++ b/posthog/test/ai/langchain/test_callbacks.py @@ -1602,8 +1602,9 @@ def test_anthropic_cache_write_and_read_tokens(mock_client): assert generation_props["$ai_input_tokens"] == 1000 assert generation_props["$ai_output_tokens"] == 50 assert generation_props["$ai_cache_creation_input_tokens"] == 800 - assert generation_props["$ai_cache_read_input_tokens"] == 0 - assert generation_props["$ai_reasoning_tokens"] == 0 + # Not reported by the fixture, so omitted rather than fabricated as 0. + assert "$ai_cache_read_input_tokens" not in generation_props + assert "$ai_reasoning_tokens" not in generation_props # Reset mock for second call mock_client.reset_mock() @@ -1637,9 +1638,9 @@ def test_anthropic_cache_write_and_read_tokens(mock_client): generation_props["$ai_input_tokens"] == 1200 ) # No provider metadata, no subtraction assert generation_props["$ai_output_tokens"] == 30 - assert generation_props["$ai_cache_creation_input_tokens"] == 0 + assert "$ai_cache_creation_input_tokens" not in generation_props assert generation_props["$ai_cache_read_input_tokens"] == 800 - assert generation_props["$ai_reasoning_tokens"] == 0 + assert "$ai_reasoning_tokens" not in generation_props def test_anthropic_provider_subtracts_cache_tokens(mock_client): @@ -1941,8 +1942,10 @@ def test_openai_cache_read_tokens(mock_client): assert generation_props["$ai_input_tokens"] == 150 # No subtraction for OpenAI assert generation_props["$ai_output_tokens"] == 40 assert generation_props["$ai_cache_read_input_tokens"] == 100 + # cache_creation is reported as an explicit 0 by the fixture, so it stays 0; + # reasoning was never reported, so it is omitted. assert generation_props["$ai_cache_creation_input_tokens"] == 0 - assert generation_props["$ai_reasoning_tokens"] == 0 + assert "$ai_reasoning_tokens" not in generation_props def test_openai_cache_creation_tokens(mock_client): @@ -1983,8 +1986,10 @@ def test_openai_cache_creation_tokens(mock_client): assert generation_props["$ai_input_tokens"] == 2000 assert generation_props["$ai_output_tokens"] == 25 assert generation_props["$ai_cache_creation_input_tokens"] == 1500 + # cache_read is reported as an explicit 0 by the fixture, so it stays 0; + # reasoning was never reported, so it is omitted. assert generation_props["$ai_cache_read_input_tokens"] == 0 - assert generation_props["$ai_reasoning_tokens"] == 0 + assert "$ai_reasoning_tokens" not in generation_props def test_combined_reasoning_and_cache_tokens(mock_client): @@ -2315,7 +2320,7 @@ def test_no_cache_read_tokens_no_subtraction(mock_client): # Input tokens should remain unchanged at 100 assert generation_props["$ai_input_tokens"] == 100 assert generation_props["$ai_output_tokens"] == 30 - assert generation_props["$ai_cache_read_input_tokens"] == 0 + assert "$ai_cache_read_input_tokens" not in generation_props def test_zero_input_tokens_with_cache_read(mock_client): @@ -2396,7 +2401,7 @@ def test_non_anthropic_cache_write_tokens_not_subtracted_from_input(mock_client) assert generation_props["$ai_input_tokens"] == 1000 assert generation_props["$ai_output_tokens"] == 20 assert generation_props["$ai_cache_creation_input_tokens"] == 800 - assert generation_props["$ai_cache_read_input_tokens"] == 0 + assert "$ai_cache_read_input_tokens" not in generation_props def test_agent_action_and_finish_imports(): diff --git a/posthog/test/ai/openai_agents/test_processor.py b/posthog/test/ai/openai_agents/test_processor.py index 1ae63f9f0..3cbc6d276 100644 --- a/posthog/test/ai/openai_agents/test_processor.py +++ b/posthog/test/ai/openai_agents/test_processor.py @@ -882,7 +882,7 @@ def test_force_flush_calls_client_flush(self, processor, mock_client): mock_client.flush.assert_called_once() def test_generation_span_with_no_usage(self, processor, mock_client, mock_span): - """Test GenerationSpanData with no usage data defaults to zero tokens.""" + """Test GenerationSpanData with no usage data omits the token counts.""" span_data = GenerationSpanData(model="gpt-4o") mock_span.span_data = span_data @@ -890,9 +890,9 @@ def test_generation_span_with_no_usage(self, processor, mock_client, mock_span): processor.on_span_end(mock_span) call_kwargs = mock_client.capture.call_args[1] - assert call_kwargs["properties"]["$ai_input_tokens"] == 0 - assert call_kwargs["properties"]["$ai_output_tokens"] == 0 - assert call_kwargs["properties"]["$ai_total_tokens"] == 0 + 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"] def test_generation_span_with_partial_usage( self, processor, mock_client, mock_span @@ -909,7 +909,9 @@ def test_generation_span_with_partial_usage( call_kwargs = mock_client.capture.call_args[1] assert call_kwargs["properties"]["$ai_input_tokens"] == 42 - assert call_kwargs["properties"]["$ai_output_tokens"] == 0 + # The output side was never reported, so it is omitted rather than 0; + # the total is the sum of the reported sides. + assert "$ai_output_tokens" not in call_kwargs["properties"] assert call_kwargs["properties"]["$ai_total_tokens"] == 42 def test_error_type_categorization_by_type_field_only( diff --git a/posthog/test/ai/test_token_reporting.py b/posthog/test/ai/test_token_reporting.py new file mode 100644 index 000000000..165dd42cf --- /dev/null +++ b/posthog/test/ai/test_token_reporting.py @@ -0,0 +1,58 @@ +"""Event token counts trace back to a provider report: absent means unknown, 0 is a report of nothing.""" + +from unittest.mock import patch + +import pytest + +from posthog.ai.types import StreamingEventData, TokenUsage +from posthog.ai.utils import capture_streaming_event + + +@pytest.fixture +def mock_client(): + with patch("posthog.client.Client") as mock_client: + mock_client.privacy_mode = False + yield mock_client + + +def _event_data(usage_stats: TokenUsage) -> StreamingEventData: + return StreamingEventData( + provider="gemini", + model="gemini-2.0-flash", + base_url="https://generativelanguage.googleapis.com", + kwargs={}, + formatted_input=[{"role": "user", "content": "hi"}], + formatted_output=[{"role": "assistant", "content": "hello"}], + usage_stats=usage_stats, + latency=0.5, + distinct_id="user-1", + trace_id="trace-1", + properties=None, + privacy_mode=False, + groups=None, + stop_reason=None, + ) + + +@pytest.mark.parametrize( + ("usage_stats", "expected"), + [ + (TokenUsage(), {}), + ( + TokenUsage(input_tokens=0, output_tokens=0), + {"$ai_input_tokens": 0, "$ai_output_tokens": 0}, + ), + (TokenUsage(input_tokens=100), {"$ai_input_tokens": 100}), + ], +) +def test_token_counts_trace_back_to_a_provider_report( + mock_client, usage_stats, expected +): + capture_streaming_event(mock_client, _event_data(usage_stats)) + + props = mock_client.capture.call_args[1]["properties"] + for key in ("$ai_input_tokens", "$ai_output_tokens"): + if key in expected: + assert props[key] == expected[key] + else: + assert key not in props From 1abf8a11811f9e252d86f35bbb82fe5eed50223b Mon Sep 17 00:00:00 2001 From: Bernat Torres Date: Mon, 31 Aug 2026 23:15:42 +0200 Subject: [PATCH 2/2] chore(ai): update public API snapshot for Optional embedding count extract_gemini_embedding_token_count now returns Optional[int], returning None when no embedding carried a token count. Regenerated with make public_api_snapshot. Generated-By: PostHog Desktop Task-Id: ec0cc27f-fc40-4994-b153-2bfe74de1edf --- references/public_api_snapshot.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index 7d74ba01c..01e2ff669 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -1025,7 +1025,7 @@ function posthog.ai.gateway.is_posthog_ai_gateway_url(base_url: Any) -> bool function posthog.ai.gateway.warn_if_posthog_ai_gateway(base_url: Any) -> None function posthog.ai.gateway.warn_if_posthog_ai_gateway_otel_attributes(attributes: Optional[Mapping[str, Any]]) -> None function posthog.ai.gemini.gemini_converter.extract_gemini_content_from_chunk(chunk: Any) -> Optional[List[FormattedContentItem]] -function posthog.ai.gemini.gemini_converter.extract_gemini_embedding_token_count(response) -> int +function posthog.ai.gemini.gemini_converter.extract_gemini_embedding_token_count(response) -> Optional[int] function posthog.ai.gemini.gemini_converter.extract_gemini_stop_reason(response: Any) -> Optional[str] function posthog.ai.gemini.gemini_converter.extract_gemini_stop_reason_from_chunk(chunk: Any) -> Optional[str] function posthog.ai.gemini.gemini_converter.extract_gemini_system_instruction(config: Any) -> Optional[str]