mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(learn/gemini): stop double-counting session tokens (#2230)
## Description
The Gemini `learn` scanner inflates every session's token totals by
double-counting.
In `_parse_messages` the per-message usage accumulation is:
```python
usage = msg.get("usageMetadata", msg.get("usage", {}))
if isinstance(usage, dict):
total_input_tokens += usage.get("promptTokenCount", 0)
total_input_tokens += usage.get("cachedContentTokenCount", 0)
total_output_tokens += usage.get("candidatesTokenCount", 0)
total_output_tokens += (
usage.get("totalTokenCount", 0) - usage.get("promptTokenCount", 0)
if usage.get("totalTokenCount")
else 0
)
```
Both additions on each side double-count, per Gemini's `usageMetadata`
semantics:
- `cachedContentTokenCount` is the cached **subset** of
`promptTokenCount`, not tokens on top of it. Adding both counts the
cached input twice.
- `totalTokenCount == promptTokenCount + candidatesTokenCount`, so
`totalTokenCount - promptTokenCount` is just `candidatesTokenCount`
again. Adding it on top of `candidatesTokenCount` counts the output
twice.
For a turn with 1000 prompt tokens (300 cached) and 500 output tokens
(`totalTokenCount` 1500), the scanner records input 1300 and output 1000
instead of 1000 / 500 — so both totals are materially inflated for any
Gemini session that carries usage metadata.
## Fix
Count the prompt as input and the candidates as output, once each:
```python
total_input_tokens += usage.get("promptTokenCount", 0)
total_output_tokens += usage.get("candidatesTokenCount", 0)
```
Closes #
## Type of Change
- [x] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)
## Changes Made
- `headroom/learn/plugins/gemini.py`: drop the `cachedContentTokenCount`
and `totalTokenCount - promptTokenCount` additions in `_parse_messages`.
- `tests/test_learn/test_gemini_scanner.py`: new test asserting the
input/output totals equal `promptTokenCount` / `candidatesTokenCount`.
- `CHANGELOG.md`: Bug Fixes entry.
## Testing
- [ ] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [ ] Manual testing performed
### Test Output
```text
$ uvx ruff@0.15.17 check headroom/learn/plugins/gemini.py tests/test_learn/test_gemini_scanner.py
All checks passed!
$ uvx mypy@1.20.2 --ignore-missing-imports headroom/learn/plugins/gemini.py
Success: no issues found in 1 source file
```
## Real Behavior Proof
- Environment: Windows 11, Python 3.12, `uvx ruff@0.15.17` / `uvx
mypy@1.20.2`. A full `pytest` OOM-kills this box (ML stack import), so I
reproduced the arithmetic with a dependency-free script and left the
full pytest to CI.
- Exact command / steps: fed a usage dict of `promptTokenCount=1000,
cachedContentTokenCount=300, candidatesTokenCount=500,
totalTokenCount=1500` through the OLD accumulation and the NEW one.
- Observed result: OLD → input 1300, output 1000 (cached and candidates
both counted twice); NEW → input 1000, output 500 (the true figures).
- Not tested: a full `learn` run over a real Gemini history; full local
`pytest` deferred to CI (OOM).
## Review Readiness
- [x] I have performed a self-review
- [x] This PR is ready for human review
## Checklist
- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [ ] New and existing unit tests pass locally with my changes
- [x] I have updated the CHANGELOG.md if applicable
## Additional Notes
The "unit tests pass locally" box is unchecked because the full suite
imports the ML stack, which I can't run here. The new test uses the
existing `GeminiScanner` harness in
`tests/test_learn/test_gemini_scanner.py` so it runs under the normal CI
pytest job; behaviour is additionally verified by the standalone proof
above.
---------
Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
This commit is contained in:
parent
7e83b8da3c
commit
29d8a5e563
2 changed files with 31 additions and 6 deletions
|
|
@ -216,14 +216,15 @@ class GeminiPlugin(LearnPlugin, ConversationScanner):
|
|||
|
||||
usage = msg.get("usageMetadata", msg.get("usage", {}))
|
||||
if isinstance(usage, dict):
|
||||
# Gemini's promptTokenCount is the FULL input token count and
|
||||
# cachedContentTokenCount is the cached SUBSET of it, so adding
|
||||
# both double-counts the cached input. Likewise totalTokenCount
|
||||
# == promptTokenCount + candidatesTokenCount, so
|
||||
# (totalTokenCount - promptTokenCount) is just candidatesTokenCount
|
||||
# again — adding it on top double-counts the output. Count the
|
||||
# prompt as input and the candidates as output, once each.
|
||||
total_input_tokens += usage.get("promptTokenCount", 0)
|
||||
total_input_tokens += usage.get("cachedContentTokenCount", 0)
|
||||
total_output_tokens += usage.get("candidatesTokenCount", 0)
|
||||
total_output_tokens += (
|
||||
usage.get("totalTokenCount", 0) - usage.get("promptTokenCount", 0)
|
||||
if usage.get("totalTokenCount")
|
||||
else 0
|
||||
)
|
||||
|
||||
if not isinstance(parts, list):
|
||||
continue
|
||||
|
|
|
|||
|
|
@ -173,6 +173,30 @@ class TestJsonSessionParsing:
|
|||
assert tc.output == "port: 8080\nhost: localhost"
|
||||
assert not tc.is_error
|
||||
|
||||
def test_token_counts_not_double_counted(self, tmp_path):
|
||||
# promptTokenCount is the full input (cachedContentTokenCount is a subset
|
||||
# of it) and totalTokenCount == prompt + candidates, so the input must be
|
||||
# promptTokenCount and the output candidatesTokenCount, each once.
|
||||
gemini_dir, _ = _setup_gemini_dir(tmp_path)
|
||||
scanner = GeminiScanner(gemini_dir=gemini_dir)
|
||||
messages = [
|
||||
{
|
||||
"role": "model",
|
||||
"parts": [{"text": "ok"}],
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": 1000,
|
||||
"cachedContentTokenCount": 300,
|
||||
"candidatesTokenCount": 500,
|
||||
"totalTokenCount": 1500,
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
session = scanner._parse_messages("s1", messages)
|
||||
|
||||
assert session.total_input_tokens == 1000
|
||||
assert session.total_output_tokens == 500
|
||||
|
||||
def test_multiple_tool_calls(self, tmp_path):
|
||||
gemini_dir, chats_dir = _setup_gemini_dir(tmp_path)
|
||||
session = _make_gemini_session(
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue