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 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]