mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-14 12:16:20 +00:00
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:
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
|
||||
Reference in new issue
Block a user