fix: accurate per-call token usage and oversized-context guards

Token accounting:
- Drain each tool round's provider stream to exhaustion before running
  tools and recursing, so the usage decorator persists exactly one
  token_usage row per LLM call, at call end. Previously every round's
  generator was abandoned mid-iteration and flushed together at request
  teardown, writing N near-identical rows stamped with the final
  round's provider counts (duplicate billing).
- Consume the Chat Completions include_usage terminal chunk (it arrives
  after finish_reason and was never read) so streamed calls record
  provider-exact token counts instead of tiktoken estimates.
- Claim provider-reported usage once per call (_last_usage_claimed) so
  a late-finalized generator can never adopt another call's counts.

Oversized-context guards:
- Enforce Responses API function_call/function_call_output pairing in
  the input builder (drop unpaired items; bypassed for store-mode
  previous_response_id chaining where calls are matched server-side).
- Hard pre-send context gate: shrink oversized tool results and refuse
  payloads that cannot fit the model's window before dispatch, so a
  hopeless request is never sent or billed.
- Cap a single tool result entering the LLM context
  (TOOL_RESULT_MAX_TOKENS, default 20000); the tool journal and
  persistence keep the full result. Applied on the resume/continuation
  path too.
- Skip the fallback attempt when the payload cannot fit the fallback
  model's context window (10% estimation slack).
- Compression: never save a compression point that does not reduce
  tokens; bound oversized verbatim fields kept after a compression
  point (COMPRESSION_RECENT_FIELD_MAX_TOKENS, default 8000).

Robustness fixes from review:
- Google parallel function calls: complete index-less ToolCalls are no
  longer merged into one another (dict arguments raised TypeError on
  +=; second call could execute with the first call's arguments).
- Trailing-frame failures after a delivered answer no longer error the
  stream or restream the whole answer from the fallback
  (_stream_reached_finish).
- In-memory compression falls back to minimal pruning when the summary
  is not smaller than the original.
- keep<=0 guard in the middle-truncation helpers (a tiny cap returned
  marker + full text).

Frontend: tooltip on the Analytics tokens stat card explaining that
agent tool loops re-send conversation context on every step.
This commit is contained in:
Alex committed 2026-07-15 21:20:39 +01:00
1 parent a7fd6d7c02
commit ca5f80995d
16 files changed
+1626 -119

No files matched your search

@@ -421,3 +421,128 @@ class TestLogToolCallStats:
}
]
svc._log_tool_call_stats(queries)
@pytest.mark.unit
class TestNegativeSavingsGuard:
"""A summary that isn't smaller than the original must never be saved
as a compression point."""
def test_growing_compression_raises(self, mock_llm, sample_conversation):
mock_llm.gen.return_value = (
"<summary>" + ("padding words that inflate the summary " * 200) + "</summary>"
)
service = CompressionService(llm=mock_llm, model_id="m")
with pytest.raises(ValueError, match="did not reduce"):
service.compress_conversation(sample_conversation, 2)
def test_growing_compression_never_saved(
self, mock_llm, mock_conversation_service, sample_conversation
):
mock_llm.gen.return_value = (
"<summary>" + ("padding words that inflate the summary " * 200) + "</summary>"
)
service = CompressionService(
llm=mock_llm,
model_id="m",
conversation_service=mock_conversation_service,
)
with pytest.raises(ValueError, match="did not reduce"):
service.compress_and_save("conv-1", sample_conversation, 2)
mock_conversation_service.update_compression_metadata.assert_not_called()
def test_shrinking_compression_still_succeeds(self, mock_llm, sample_conversation):
# Default fixture summary is tiny relative to the conversation.
service = CompressionService(llm=mock_llm, model_id="m")
metadata = service.compress_conversation(sample_conversation, 2)
assert metadata.compressed_token_count < metadata.original_token_count
@pytest.mark.unit
class TestBoundRecentQueries:
"""Oversized verbatim fields in the post-compression tail are capped."""
def _compressed_conversation(self, recent_query):
return {
"queries": [
{"prompt": "old", "response": "old answer"},
recent_query,
],
"compression_metadata": {
"is_compressed": True,
"compression_points": [
{
"compressed_summary": "summary",
"query_index": 0,
"compressed_token_count": 10,
"original_token_count": 100,
}
],
},
}
def test_oversized_response_is_trimmed_without_mutating_original(
self, mock_llm, monkeypatch
):
monkeypatch.setattr(
"application.api.answer.services.compression.service.settings.COMPRESSION_RECENT_FIELD_MAX_TOKENS",
50,
raising=False,
)
big = "word " * 2000
recent = {"prompt": "q", "response": big}
conversation = self._compressed_conversation(recent)
service = CompressionService(llm=mock_llm, model_id="m")
summary, queries = service.get_compressed_context(conversation)
assert summary == "summary"
assert len(queries) == 1
assert "trimmed to fit context" in queries[0]["response"]
assert len(queries[0]["response"]) < len(big)
# Original conversation dict untouched.
assert recent["response"] == big
def test_oversized_tool_result_is_trimmed(self, mock_llm, monkeypatch):
monkeypatch.setattr(
"application.api.answer.services.compression.service.settings.COMPRESSION_RECENT_FIELD_MAX_TOKENS",
50,
raising=False,
)
big = "data " * 2000
recent = {
"prompt": "q",
"response": "short",
"tool_calls": [{"tool_name": "t", "result": big}],
}
conversation = self._compressed_conversation(recent)
service = CompressionService(llm=mock_llm, model_id="m")
_summary, queries = service.get_compressed_context(conversation)
assert "trimmed to fit context" in queries[0]["tool_calls"][0]["result"]
assert recent["tool_calls"][0]["result"] == big
def test_small_fields_pass_through_unchanged(self, mock_llm):
recent = {"prompt": "q", "response": "short answer"}
conversation = self._compressed_conversation(recent)
service = CompressionService(llm=mock_llm, model_id="m")
_summary, queries = service.get_compressed_context(conversation)
assert queries[0] is recent
def test_zero_cap_disables_bounding(self, mock_llm, monkeypatch):
monkeypatch.setattr(
"application.api.answer.services.compression.service.settings.COMPRESSION_RECENT_FIELD_MAX_TOKENS",
0,
raising=False,
)
big = "word " * 2000
recent = {"prompt": "q", "response": big}
conversation = self._compressed_conversation(recent)
service = CompressionService(llm=mock_llm, model_id="m")
_summary, queries = service.get_compressed_context(conversation)
assert queries[0] is recent