Skip to content

Capture real token counts from backends and use for compaction#33

Merged
antoinezambelli merged 2 commits into
mainfrom
az/token-llama
Mar 27, 2026
Merged

Capture real token counts from backends and use for compaction#33
antoinezambelli merged 2 commits into
mainfrom
az/token-llama

Conversation

@antoinezambelli

Copy link
Copy Markdown
Owner

Capture real token counts from backends and use for compaction

Summary

  • LlamafileClient and OllamaClient now record actual token usage (TokenUsage dataclass) from every LLM response into a last_usage dict keyed by slot ID
  • ContextManager uses real token counts (via update_token_count()) instead of the chars/4 heuristic when available
  • run_inference() syncs token counts from the client to the context manager after every LLM call
  • Streaming path requests stream_options: {include_usage: true} and breaks on [DONE] instead of finish_reason to capture the usage chunk that arrives between the two

Why

The chars/4 heuristic significantly underestimates actual token usage — especially for reasoning models where <think> tokens consume KV cache but aren't persisted in message content. In forge-code eval runs, 14B Q8 reasoning models StreamError at 6-9K estimated tokens against a 12K budget because compaction never fires — the heuristic thinks there's headroom when KV cache is actually full.

With real counts, compaction and context warnings trigger based on actual server-reported usage.

Changes

File Change
clients/base.py New TokenUsage frozen dataclass
clients/llamafile.py last_usage dict, _record_usage(), stream_options, break on [DONE]
clients/ollama.py last_usage dict, _record_usage() (maps prompt_eval_count/eval_count)
context/manager.py update_token_count(), estimate_tokens() returns real count when available
core/inference.py _sync_token_count() called after every send
__init__.py Export TokenUsage

Design decisions

  • last_usage is a dict[int, TokenUsage] keyed by slot ID — supports future sub-agent work (issues Sub-agent support: synchronous swap scheduling #27-Sub-agent support: slot pool #29) where concurrent agents share a client with multiple slots. Default slot 0 for single-slot use.
  • Client-side state rather than return type change — avoids breaking the LLMResponse union type. _sync_token_count uses getattr so clients without last_usage (e.g. AnthropicClient) are silently skipped.
  • Dropped llamafile 0.9.x compat — streaming now breaks on [DONE] instead of finish_reason to capture the usage chunk. Old llamafile binaries that don't send [DONE] are no longer supported.

Test plan

  • 677 unit tests pass
  • Integration verified against live llama-server (both streaming and non-streaming)
  • Integration verified against live Ollama (both streaming and non-streaming)
  • forge-code eval_runner smoke test passed on llamaserver

Closes #32

antoinezambelli and others added 2 commits March 26, 2026 23:53
LlamafileClient now records TokenUsage (prompt, completion, total) from
every llama-server response. ContextManager uses these real counts
instead of chars/4 heuristic when available, so compaction and context
warnings fire based on actual KV cache pressure — critical for reasoning
models where think tokens consume context invisibly.

Closes #32

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Same pattern as LlamafileClient — captures prompt_eval_count and
eval_count from Ollama responses (both streaming and non-streaming)
into last_usage dict. ContextManager picks these up via the existing
_sync_token_count path in run_inference.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
@antoinezambelli
antoinezambelli merged commit 1280237 into main Mar 27, 2026
2 checks passed
@antoinezambelli
antoinezambelli deleted the az/token-llama branch March 27, 2026 04:54
isgallagher pushed a commit to isgallagher/forge-guardrails that referenced this pull request May 29, 2026
…nezambelli#33)

* Capture real token counts from llama-server and use for compaction

LlamafileClient now records TokenUsage (prompt, completion, total) from
every llama-server response. ContextManager uses these real counts
instead of chars/4 heuristic when available, so compaction and context
warnings fire based on actual KV cache pressure — critical for reasoning
models where think tokens consume context invisibly.

Closes antoinezambelli#32

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Add token usage capture to OllamaClient

Same pattern as LlamafileClient — captures prompt_eval_count and
eval_count from Ollama responses (both streaming and non-streaming)
into last_usage dict. ContextManager picks these up via the existing
_sync_token_count path in run_inference.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Token usage from llama-server responses is discarded

1 participant