Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions .sampo/changesets/openai-agents-cache-reasoning-tokens.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
pypi/posthog: patch
---

Report cached input tokens from OpenAI Agents generation and response spans, which nest them under the usage details fields, and declare them inclusive.
52 changes: 49 additions & 3 deletions posthog/ai/openai_agents/processor.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,27 @@ def _parse_iso_timestamp(iso_str: Optional[str]) -> Optional[float]:
return None


def _usage_detail(usage: Any, details_keys: tuple, key: str) -> Optional[int]:
"""Read a token count nested under a usage details field.

Agents span usage is a dict (generation spans) or an OpenAI usage object
(response spans), with details under ``input_tokens_details``/
``output_tokens_details`` (Responses) or ``prompt_tokens_details``/
``completion_tokens_details`` (Chat Completions).
"""

def get(container: Any, name: str) -> Any:
if isinstance(container, dict):
return container.get(name)
return getattr(container, name, None)

for details_key in details_keys:
value = get(get(usage, details_key), key)
if isinstance(value, int):
return value
return None


class PostHogTracingProcessor(TracingProcessor):
"""
A tracing processor that sends OpenAI Agents SDK traces to PostHog.
Expand Down Expand Up @@ -528,11 +549,27 @@ def _handle_generation_span(
),
}

# Add optional token fields if present
# Add optional token fields if present. The Agents SDK nests cached
# tokens under the usage details fields; flat keys are kept for custom
# models. Nested reasoning is not forwarded: compatible providers count
# it inside the output total, and the cost backend would add it again.
if usage.get("reasoning_tokens"):
properties["$ai_reasoning_tokens"] = usage["reasoning_tokens"]
if usage.get("cache_read_input_tokens"):
properties["$ai_cache_read_input_tokens"] = usage["cache_read_input_tokens"]
cache_read_tokens = usage.get("cache_read_input_tokens")
if not cache_read_tokens:
cache_read_tokens = _usage_detail(
usage,
("input_tokens_details", "prompt_tokens_details"),
"cached_tokens",
)
if cache_read_tokens:
# OpenAI-shaped usage counts cached tokens inside input_tokens.
# Say so explicitly: the cost backend otherwise infers exclusive
# reporting from Claude model names served through compatible
# endpoints. A caller-supplied value still takes precedence.
properties["$ai_cache_reporting_exclusive"] = False
if cache_read_tokens:
properties["$ai_cache_read_input_tokens"] = cache_read_tokens
Comment thread
marandaneto marked this conversation as resolved.
if usage.get("cache_creation_input_tokens"):
properties["$ai_cache_creation_input_tokens"] = usage[
"cache_creation_input_tokens"
Expand Down Expand Up @@ -708,6 +745,15 @@ def _handle_response_span(
if total_cost_usd is not None:
properties["$ai_total_cost_usd"] = total_cost_usd

if usage:
cache_read_tokens = _usage_detail(
usage, ("input_tokens_details",), "cached_tokens"
)
if cache_read_tokens:
properties["$ai_cache_read_input_tokens"] = cache_read_tokens
# Responses input_tokens include cached tokens.
properties["$ai_cache_reporting_exclusive"] = False

# Extract output content from response
if response:
output_items = getattr(response, "output", None)
Expand Down
157 changes: 157 additions & 0 deletions posthog/test/ai/openai_agents/test_processor.py
Original file line number Diff line number Diff line change
Expand Up @@ -288,6 +288,163 @@ def test_generation_span_with_reasoning_tokens(
call_kwargs = mock_client.capture.call_args[1]
assert call_kwargs["properties"]["$ai_reasoning_tokens"] == 400

def test_generation_span_reads_nested_cache_tokens(
self, processor, mock_client, mock_span
):
"""The Agents SDK nests cached and reasoning tokens under the usage
details fields of generation span usage."""
from agents.usage import Usage, model_usage_to_span_usage
from openai.types.responses.response_usage import (
InputTokensDetails,
OutputTokensDetails,
)

usage = Usage(
requests=1,
input_tokens=100,
output_tokens=50,
total_tokens=150,
input_tokens_details=InputTokensDetails(
cached_tokens=60, cache_write_tokens=0
),
output_tokens_details=OutputTokensDetails(reasoning_tokens=40),
)
mock_span.span_data = GenerationSpanData(
model="gpt-5-mini", usage=model_usage_to_span_usage(usage)
)

processor.on_span_start(mock_span)
processor.on_span_end(mock_span)

props = mock_client.capture.call_args[1]["properties"]
assert props["$ai_input_tokens"] == 100
assert props["$ai_cache_read_input_tokens"] == 60
assert props["$ai_output_tokens"] == 50
# Inclusive in output_tokens; forwarding it would be counted twice.
assert "$ai_reasoning_tokens" not in props
assert props["$ai_cache_reporting_exclusive"] is False

def test_response_span_reads_cache_tokens(self, processor, mock_client, mock_span):
"""Response spans carry the OpenAI Responses usage object, with cached
and reasoning tokens in its details fields."""
from openai.types.responses import Response

response = Response.model_validate(
{
"id": "resp_123",
"object": "response",
"created_at": 0,
"model": "gpt-5-mini",
"status": "completed",
"parallel_tool_calls": True,
"tool_choice": "auto",
"tools": [],
"output": [],
"usage": {
"input_tokens": 100,
"output_tokens": 50,
"total_tokens": 150,
"input_tokens_details": {
"cached_tokens": 60,
"cache_write_tokens": 0,
},
"output_tokens_details": {"reasoning_tokens": 40},
},
}
)
mock_span.span_data = ResponseSpanData(response=response, input="Hi")

processor.on_span_start(mock_span)
processor.on_span_end(mock_span)

props = mock_client.capture.call_args[1]["properties"]
assert props["$ai_input_tokens"] == 100
assert props["$ai_cache_read_input_tokens"] == 60
assert props["$ai_output_tokens"] == 50
# Inclusive in output_tokens; forwarding it would be counted twice.
assert "$ai_reasoning_tokens" not in props
assert props["$ai_cache_reporting_exclusive"] is False

def test_cached_compatible_claude_usage_is_declared_inclusive(self):
"""A Claude model behind an OpenAI-compatible endpoint reports inclusive
cached tokens. Without an explicit convention, the cost backend infers
exclusive reporting from the model name and overcounts input cost."""
from agents import set_trace_processors
from agents.tracing import generation_span, trace

client = MagicMock(privacy_mode=False, enable_full_ai_capture=False)
set_trace_processors([])
processor = instrument(client, distinct_id="review-user")
try:
with trace("compatible-model"):
with generation_span(
model="anthropic/claude-sonnet-4.6",
usage={
"input_tokens": 10339,
"output_tokens": 60,
"input_tokens_details": {"cached_tokens": 10318},
},
):
pass
properties = next(
call.kwargs["properties"]
for call in client.capture.call_args_list
if call.kwargs["event"] == "$ai_generation"
)
assert properties["$ai_cache_read_input_tokens"] == 10318
assert properties.get("$ai_cache_reporting_exclusive") is False
finally:
set_trace_processors([])
processor.shutdown()

def test_flat_cache_keys_keep_their_own_convention(
self, processor, mock_client, mock_span
):
"""Flat cache keys come from custom models with their own convention, so
no inclusive/exclusive flag is declared for them."""
mock_span.span_data = GenerationSpanData(
model="custom-model",
usage={
"input_tokens": 100,
"output_tokens": 50,
"cache_read_input_tokens": 60,
},
)

processor.on_span_start(mock_span)
processor.on_span_end(mock_span)

props = mock_client.capture.call_args[1]["properties"]
assert props["$ai_cache_read_input_tokens"] == 60
assert "$ai_cache_reporting_exclusive" not in props

def test_caller_cache_reporting_convention_takes_precedence(
self, mock_client, mock_span
):
"""A processor-wide $ai_cache_reporting_exclusive overrides the default
declared for nested OpenAI-shaped usage."""
processor = PostHogTracingProcessor(
client=mock_client,
distinct_id="test-user",
privacy_mode=False,
properties={"$ai_cache_reporting_exclusive": True},
)
mock_span.span_data = GenerationSpanData(
model="gpt-5-mini",
usage={
"input_tokens": 100,
"output_tokens": 50,
"input_tokens_details": {"cached_tokens": 60},
},
)

processor.on_span_start(mock_span)
processor.on_span_end(mock_span)

props = mock_client.capture.call_args[1]["properties"]
assert props["$ai_cache_read_input_tokens"] == 60
assert props["$ai_cache_reporting_exclusive"] is True

def test_function_span_mapping(self, processor, mock_client, mock_span):
"""Test FunctionSpanData maps to $ai_span event with type=tool."""
span_data = FunctionSpanData(
Expand Down
Loading