Skip to content

fix(llm): keep the cached-token count the server reports - #4032

Closed
kovtcharov wants to merge 3 commits into
amd:mainfrom
kovtcharov:claude/report-cached-tokens
Closed

kovtcharov wants to merge 3 commits into
amd:mainfrom
kovtcharov:claude/report-cached-tokens

Conversation

@kovtcharov

Copy link
Copy Markdown
Contributor

GAIA prices every input token as if it were fresh. The provider copied four fields out of each completion's usage and dropped the count of tokens the server served from its cache, so the cost and token totals users see are upper bounds with no way to tell how far off they are. A benchmark measurement put the cache hit rate at roughly 95% within a task, which makes the gap large. The count is now recorded per step and totalled per turn.

Test plan

  • python -m pytest tests/unit/test_agent_token_usage.py passes. On main the 6 new tests fail.
  • A completion with prompt_tokens_details.cached_tokens, one with a flat cached_tokens, and one with neither are all handled.
  • gaia stats and any cost display still work for a backend that reports no cache (reads as 0, not as unknown).
🔍 Technical details

LemonadeProvider._capture_usage() now holds the usage-capture logic that was inline, and adds cached_tokens. get_performance_stats() already copies _last_usage wholesale, so the count reaches each step's performance_stats entry and from there the conversation. _sum_cached_tokens(conversation) totals it alongside the existing _sum_conversation_tokens. #3739 adds a function of the same name for the TUI's session cost line; if that lands first, the two should be merged into one.

Every token figure GAIA reports treats a cached input token as a fresh one,
because the provider copied four fields out of the completion's usage and
dropped prompt_tokens_details.cached_tokens. A benchmark measurement put the
cache hit rate at about 95% within a task, so the cost shown for a run is an
upper bound with no way to tell how far off it is.

The count now rides with the rest of the usage, per step, and
_sum_cached_tokens totals it for a turn. Both spellings are read: OpenAI's
prompt_tokens_details.cached_tokens and a flat cached_tokens. A backend that
does not cache reports zero.
@github-actions github-actions Bot added llm LLM backend changes tests Test changes performance Performance-critical changes agents labels Sep 19, 2026
@github-actions

Copy link
Copy Markdown
Contributor

Request changes

The PR captures the cached-token count from each completion, which is the right half of the fix — but nothing ever reads it back out. The per-turn total the description promises is written as a helper, tested, and then never called, so gaia stats, the agent's returned token totals, and every cost figure downstream still show exactly what they showed before this PR. From a user's point of view the cost numbers are unchanged upper bounds.

To finish it, the new per-turn sum needs to be called where the other token totals are computed at the end of a query, and its result carried in the returned totals alongside input and output tokens — otherwise the count stops at the provider and never reaches anyone.

Two smaller things worth a look:

  • The provider-side capture is less defensive than its agent-side counterpart: if the server ever reports the count as something non-numeric rather than as a number or absent, the capture raises and takes the chat call down with it. The agent-side helper already guards against exactly that; the two should agree.
  • Streaming turns still have no cached count at all — the streaming path never records usage, so those turns fall back to the server's polled stats, which the code itself notes only counts the uncached part. Pre-existing, but it means the improvement only lands on non-streaming turns; worth saying so in the description if that's intended.

The test coverage for what is implemented is good, and the two response shapes plus the no-cache case are all handled.

Real-world evidence

N/A — no evidence bundle was produced for this run, and no CLI/API output is attached to the PR. The three test-plan boxes are all unchecked, including the gaia stats one that would have caught the missing wiring. I could not run the suite here either (pytest is not installed in this reviewer environment), so this verdict rests on static review alone. Running the third test-plan item on a real backend is the check that matters most: with the current diff, gaia stats output is byte-identical before and after.

🔍 Technical details

Issues

🟡 _sum_cached_tokens is dead code — the per-turn total never reaches the result (src/gaia/agents/base/agent.py:465)

_sum_cached_tokens has exactly two references in the tree: its definition and tests/unit/test_agent_token_usage.py. The per-step half works — _capture_usage → get_performance_stats() → perf_stats → the type: "stats" conversation entry (agent.py:6938) — but _process_query_impl only calls _sum_conversation_tokens (agent.py:7480) and the result dict it builds at agent.py:7509 carries input_tokens / output_tokens / total_tokens and no cached count. The description's "totalled per turn" is not true of this diff.

        total_input_tokens, total_output_tokens = _sum_conversation_tokens(
            conversation, self._tool_reported_usage
        )
        total_cached_tokens = _sum_cached_tokens(conversation)

…then add "cached_tokens": total_cached_tokens, to the result dict next to "total_tokens". Worth checking the same for the pre-answer computation at agent.py:7394 if that total is surfaced anywhere.

🟢 _capture_usage can raise on a malformed count (src/gaia/llm/providers/lemonade.py:532)

int(cached or 0) propagates a ValueError out of chat() if the server ever sends a non-numeric value, while _sum_cached_tokens (agent.py:480) catches (TypeError, ValueError) for the same field. test_a_non_numeric_count_does_not_raise passes None, which is the absent case that or 0 already covers — it does not exercise the name it claims. Either rename the test to match what it asserts, or make the two sides agree by coercing defensively here.

🟢 Streaming turns still report no cached prefix (src/gaia/llm/providers/lemonade.py:568)

_handle_stream never assigns _last_usage, so get_performance_stats() falls through to self._backend.get_stats() — which the comment at lemonade.py:546 says "counts only the uncached part". Pre-existing and out of scope to fix here, but it bounds the change: only non-streaming turns get the new count, which is worth one line in the PR description.

Strengths

  • _capture_usage handles both the nested prompt_tokens_details.cached_tokens and a flat cached_tokens, preferring the OpenAI shape only when it is actually present rather than when the parent key exists — that ordering is easy to get wrong.
  • Extracting the inline capture into a named method makes it unit-testable without an HTTP layer, and the tests take advantage of that instead of mocking a server.
  • The docstrings say why (cached tokens are priced differently; absent reads as zero, not unknown) rather than restating the code.

@github-actions

Copy link
Copy Markdown
Contributor

🟡 The previous review's central concern is still open: _sum_cached_tokens is dead code. The function was added and thoroughly tested, but it is never called — _process_query_impl still only invokes _sum_conversation_tokens, and the result dict carries no cached_tokens key. From a user's perspective the cost figures are unchanged upper bounds, same as before this PR.

The fix is two lines — call the function and include its result alongside the other totals.

🔍 Technical details

src/gaia/agents/base/agent.py:7480 — add the call beside _sum_conversation_tokens:

        total_input_tokens, total_output_tokens = _sum_conversation_tokens(
            conversation, self._tool_reported_usage
        )
        total_cached_tokens = _sum_cached_tokens(conversation)

Then add to the result dict (~line 7522):

            "cached_tokens": total_cached_tokens,

_capture_usage and its tests are correct and do land the per-step data in performance_stats; the only missing step is reading it back out at the end of the turn.

@kovtcharov-amd

Copy link
Copy Markdown
Collaborator

Refreshed this branch against current main so it can merge again — the branch had fallen far enough behind that GitHub could no longer compute a merge.

The update was a clean merge: nothing was resolved by hand, and no commit on this branch was rewritten. CI now runs against an up-to-date base.

@kovtcharov-amd

Copy link
Copy Markdown
Collaborator

Not merging this — main already does everything the PR sets out to do, and a bit more, so landing it now would take capability away rather than add it.

The PR description called this: "#3739 adds a function of the same name for the TUI's session cost line; if that lands first, the two should be merged into one." #3739 landed first, and it shipped the whole thing — the server-reported cached count is captured per call, summed per turn, and actually displayed. The two review rounds asking for the missing wiring were answered by that PR, not this one.

Recommend closing as superseded.

🔍 Technical details

Verified against origin/main @ 065bb810b, not against the PR's stale base.

Already on main, wired end to end:

  • _sum_cached_tokens(conversation) — src/gaia/agents/base/agent.py:553, and called at agent.py:8179, feeding print_final_answer(cached_tokens=...). This is exactly the dead-code finding from both review rounds, already closed.
  • LemonadeProvider._capture_usage(usage, timings) — src/gaia/llm/providers/lemonade.py:698, reading prompt_tokens_details.cached_tokens.
  • TestSumCachedTokens in tests/unit/test_agent_token_usage.py — 6 cases including malformed/negative/None counts; plus the capture-side cases in test_lemonade_cloud_models.py:573-595.

python -m pytest tests/unit/test_agent_token_usage.py tests/unit/test_lemonade_cloud_models.py → 114 passed on main as-is.

Three ways merging this branch would regress main:

  1. Streaming loses its count again. Main calls _capture_usage from _handle_stream (lemonade.py:774) as well as the non-streaming path. This branch's single-argument _capture_usage(response) only runs in chat() — that is the third bot finding, fixed on main, undone here.
  2. reasoning_tokens is dropped. Main's capture also records _reasoning_tokens(usage) (lemonade.py:57, 719). This branch's rewrite of the method has no equivalent.
  3. "Reported 0" and "reported nothing" stop being distinguishable. Main omits cached_tokens when the backend sent none and keeps a genuine 0 as a measurement. This branch's int(cached or 0) invents a 0 for a silent backend — and its test_lemonade_cloud_models.py change asserts that invented 0, contradicting main's assert "cached_tokens" not in captured.

The only thing here with no counterpart on main is the flat top-level cached_tokens fallback. No backend in the tree emits that shape — OpenAI and Fireworks both use the nested form — so it is speculative defensiveness, not a gap worth reopening for.

No eval run: the change is token-usage bookkeeping. It touches no system prompt, tool schema, error classification, default model, or tool-call parsing, so none of the CLAUDE.md eval triggers apply. Moot regardless, since the recommendation is to close.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

agents llm LLM backend changes performance Performance-critical changes tests Test changes

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants