Skip to content

Commit d9bc612

Browse files
authored
fix(openai): record cached and reasoning tokens on the LLM metrics counters (#2758)
The OpenAI-compatible provider extracts cached_tokens and thoughts_tokens on both call paths and hands them to TokenUsage, but never passes them to metrics.record_llm_call, which accepts and buckets both. Two separate effects: - Reasoning tokens reach no counter at all. #2378 made output_tokens visible-only by subtracting thoughts_tokens directly above the record_llm_call, so the reasoning half of the billed output was removed from the metrics path rather than moved onto llm_tokens_thoughts. Before #2378 those tokens were still counted inside output_tokens. - cached_input_tokens has read 0 for every OpenAI-compatible provider since the counter was added; only gemini_llm passes it. Pass both kwargs at the two call sites that parse a usage object. The fallback path (no usage) and the Ollama native path (no reasoning or cached fields) are unchanged. Invariant: recorded output_tokens + recorded thoughts_tokens equals the provider's completion_tokens, so every billed token lands on exactly one counter. The new tests assert on the collector itself; the existing ones patch it without asserting, which is why this went unnoticed.
1 parent 1fe43ec commit d9bc612

2 files changed

Lines changed: 94 additions & 2 deletions

File tree

hindsight-api-slim/hindsight_api/engine/providers/openai_compatible_llm.py

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -959,7 +959,9 @@ async def call(
959959
output_tokens = max(0, output_tokens - thoughts_tokens)
960960
total_tokens = max(0, total_tokens - thoughts_tokens)
961961

962-
# Record LLM metrics
962+
# Record LLM metrics. ``output_tokens`` is visible-only by now, so
963+
# ``thoughts_tokens`` has to be recorded alongside it or the reasoning
964+
# half of the billed output reaches no counter at all.
963965
metrics = get_metrics_collector()
964966
metrics.record_llm_call(
965967
provider=self.provider,
@@ -969,6 +971,8 @@ async def call(
969971
input_tokens=input_tokens,
970972
output_tokens=output_tokens,
971973
success=True,
974+
cached_input_tokens=cached_tokens,
975+
thoughts_tokens=thoughts_tokens,
972976
)
973977

974978
# Record trace span
@@ -1270,6 +1274,8 @@ async def call_with_tools(
12701274
if thoughts_tokens:
12711275
output_tokens = max(0, output_tokens - thoughts_tokens)
12721276

1277+
# See ``call()``: record the reasoning and cached counts too, so no
1278+
# billed token is dropped from the metrics counters.
12731279
metrics = get_metrics_collector()
12741280
metrics.record_llm_call(
12751281
provider=self.provider,
@@ -1279,6 +1285,8 @@ async def call_with_tools(
12791285
input_tokens=input_tokens,
12801286
output_tokens=output_tokens,
12811287
success=True,
1288+
cached_input_tokens=cached_tokens,
1289+
thoughts_tokens=thoughts_tokens,
12821290
)
12831291

12841292
# Record OpenTelemetry span

hindsight-api-slim/tests/test_token_usage_cached_thoughts.py

Lines changed: 85 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
from __future__ import annotations
1818

1919
from types import SimpleNamespace
20-
from unittest.mock import AsyncMock, patch
20+
from unittest.mock import AsyncMock, MagicMock, patch
2121

2222
import pytest
2323
from pydantic import BaseModel
@@ -27,6 +27,7 @@
2727
from hindsight_api.engine.reflect.models import StructuredOutputResult, TokenUsageSummary
2828
from hindsight_api.engine.response_models import LLMToolCallResult, TokenUsage
2929
from hindsight_api.extensions.operation_validator import RetainResult
30+
from hindsight_api.metrics import MetricsCollector
3031

3132

3233
class _OkModel(BaseModel):
@@ -293,3 +294,86 @@ async def test_openai_compatible_output_tokens_exclude_thoughts_like_gemini():
293294
assert token_usage.total_tokens == token_usage.input_tokens + token_usage.output_tokens
294295
# The reasoning tokens live in exactly one field, not both.
295296
assert token_usage.output_tokens + token_usage.thoughts_tokens == 83
297+
298+
299+
def _recorded_llm_call(collector):
300+
"""The kwargs of the single record_llm_call the provider made."""
301+
assert collector.record_llm_call.call_count == 1
302+
return collector.record_llm_call.call_args.kwargs
303+
304+
305+
@pytest.mark.asyncio
306+
async def test_openai_compatible_call_records_cached_and_thoughts_on_metrics():
307+
"""call() must hand cached/thoughts to the metrics collector, not only to
308+
TokenUsage. The collector already accepts both kwargs and keeps a counter for
309+
each; the provider simply never passed them, so hindsight.llm.tokens.thoughts
310+
and hindsight.llm.tokens.cached_input stayed empty for every OpenAI-compatible
311+
reasoning model."""
312+
llm = _openai_llm()
313+
llm._client.chat.completions.create = AsyncMock(return_value=_response(usage=_usage(cached=200, reasoning=80)))
314+
collector = MagicMock(spec=MetricsCollector)
315+
with patch(
316+
"hindsight_api.engine.providers.openai_compatible_llm.get_metrics_collector",
317+
return_value=collector,
318+
):
319+
await llm.call(
320+
messages=[{"role": "user", "content": "Return whether this worked."}],
321+
response_format=_OkModel,
322+
max_retries=0,
323+
return_usage=True,
324+
)
325+
recorded = _recorded_llm_call(collector)
326+
assert recorded["cached_input_tokens"] == 200
327+
assert recorded["thoughts_tokens"] == 80
328+
329+
330+
@pytest.mark.asyncio
331+
async def test_openai_compatible_call_with_tools_records_cached_and_thoughts_on_metrics():
332+
"""call_with_tools() extracts both counts for its own return value, so it must
333+
record them too — the tool-calling path is where the agentic loop spends most
334+
of its reasoning tokens."""
335+
llm = _openai_llm()
336+
llm._client.chat.completions.create = AsyncMock(
337+
return_value=_response(usage=_usage(cached=64, reasoning=33), content="done")
338+
)
339+
collector = MagicMock(spec=MetricsCollector)
340+
with patch(
341+
"hindsight_api.engine.providers.openai_compatible_llm.get_metrics_collector",
342+
return_value=collector,
343+
):
344+
await llm.call_with_tools(messages=[{"role": "user", "content": "hi"}], tools=[], max_retries=0)
345+
recorded = _recorded_llm_call(collector)
346+
assert recorded["cached_input_tokens"] == 64
347+
assert recorded["thoughts_tokens"] == 33
348+
349+
350+
@pytest.mark.asyncio
351+
async def test_openai_compatible_metrics_account_for_every_billed_output_token():
352+
"""Accounting invariant: every token the provider bills as output lands on
353+
exactly one counter, never zero and never two.
354+
355+
output_tokens is made visible-only just above the record_llm_call (reasoning
356+
subtracted so it doesn't double-count against thoughts_tokens). That subtraction
357+
only holds the books straight if thoughts_tokens is recorded as well: recorded
358+
output + recorded thoughts must add back up to the provider's completion_tokens.
359+
"""
360+
completion, reasoning = 83, 64
361+
llm = _openai_llm()
362+
llm._client.chat.completions.create = AsyncMock(
363+
return_value=_response(usage=_usage(reasoning=reasoning, prompt=20, completion=completion, total=103))
364+
)
365+
collector = MagicMock(spec=MetricsCollector)
366+
with patch(
367+
"hindsight_api.engine.providers.openai_compatible_llm.get_metrics_collector",
368+
return_value=collector,
369+
):
370+
await llm.call(
371+
messages=[{"role": "user", "content": "17*23?"}],
372+
response_format=_OkModel,
373+
max_retries=0,
374+
return_usage=True,
375+
)
376+
recorded = _recorded_llm_call(collector)
377+
assert recorded["output_tokens"] == completion - reasoning # visible-only
378+
assert recorded["thoughts_tokens"] == reasoning
379+
assert recorded["output_tokens"] + recorded["thoughts_tokens"] == completion

0 commit comments

Comments
 (0)