mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-14 16:16:48 +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
@@ -811,3 +811,124 @@ class TestRespondingProviderTracking:
|
||||
primary = _Google(fail_at=0, backup_models=["backup-model"])
|
||||
primary.gen(**CALL_ARGS)
|
||||
assert primary._responding_provider == "openai"
|
||||
|
||||
|
||||
# Tests — fallback payload size gate
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
class TestFallbackPayloadSizeGate:
|
||||
"""A payload that cannot fit the fallback's context window must skip the
|
||||
fallback attempt (it would be a guaranteed second rejection) and
|
||||
propagate the primary's error instead."""
|
||||
|
||||
def _primary_with_backup(self, patch_model_utils, backup):
|
||||
patch_model_utils(
|
||||
get_provider=lambda mid, **_kw: "openai",
|
||||
get_api_key=lambda p: "k",
|
||||
create_llm=lambda type, **kw: backup,
|
||||
)
|
||||
return FakeLLM(fail_at=0, backup_models=["backup-model"])
|
||||
|
||||
BIG_ARGS = dict(
|
||||
model="test-model",
|
||||
messages=[{"role": "user", "content": "word " * 300}],
|
||||
)
|
||||
|
||||
def test_gen_skips_fallback_when_payload_cannot_fit(
|
||||
self, monkeypatch, patch_model_utils
|
||||
):
|
||||
backup = FakeLLM(responses=["backup ok"])
|
||||
primary = self._primary_with_backup(patch_model_utils, backup)
|
||||
monkeypatch.setattr(
|
||||
"application.core.model_utils.get_token_limit",
|
||||
lambda mid, user_id=None: 10,
|
||||
)
|
||||
with pytest.raises(RuntimeError, match="primary model unavailable"):
|
||||
primary.gen(**self.BIG_ARGS)
|
||||
assert backup.gen_called is False
|
||||
|
||||
def test_stream_skips_fallback_when_payload_cannot_fit(
|
||||
self, monkeypatch, patch_model_utils
|
||||
):
|
||||
backup = FakeLLM(stream_chunks=["backup chunk"])
|
||||
primary = self._primary_with_backup(patch_model_utils, backup)
|
||||
monkeypatch.setattr(
|
||||
"application.core.model_utils.get_token_limit",
|
||||
lambda mid, user_id=None: 10,
|
||||
)
|
||||
with pytest.raises(RuntimeError, match="mid-stream failure"):
|
||||
list(primary.gen_stream(**self.BIG_ARGS))
|
||||
assert backup.gen_stream_called is False
|
||||
|
||||
def test_fallback_proceeds_when_payload_fits(
|
||||
self, monkeypatch, patch_model_utils
|
||||
):
|
||||
backup = FakeLLM(responses=["backup ok"])
|
||||
primary = self._primary_with_backup(patch_model_utils, backup)
|
||||
monkeypatch.setattr(
|
||||
"application.core.model_utils.get_token_limit",
|
||||
lambda mid, user_id=None: 100000,
|
||||
)
|
||||
assert primary.gen(**self.BIG_ARGS) == "backup ok"
|
||||
assert backup.gen_called is True
|
||||
|
||||
def test_estimation_failure_never_blocks_fallback(
|
||||
self, monkeypatch, patch_model_utils
|
||||
):
|
||||
backup = FakeLLM(responses=["backup ok"])
|
||||
primary = self._primary_with_backup(patch_model_utils, backup)
|
||||
|
||||
def boom(*a, **kw):
|
||||
raise ValueError("estimator broken")
|
||||
|
||||
monkeypatch.setattr("application.usage._count_prompt_tokens", boom)
|
||||
assert primary.gen(**self.BIG_ARGS) == "backup ok"
|
||||
assert backup.gen_called is True
|
||||
|
||||
|
||||
# Tests — no fallback restream after the primary delivered its finish signal
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
class TestNoRestreamAfterFinish:
|
||||
"""A trailing-frame failure (between the finish chunk and stream end)
|
||||
must NOT restream the already-delivered answer from the fallback."""
|
||||
|
||||
class FinishThenFailLLM(FakeLLM):
|
||||
def _raw_gen_stream(self, baseself, model, messages, stream, tools=None, **kwargs):
|
||||
self.gen_stream_called = True
|
||||
yield "the full answer"
|
||||
# Provider marked the stream finished (finish_reason arrived)...
|
||||
self._stream_reached_finish = True
|
||||
# ...then the trailing usage/[DONE] frame dies.
|
||||
raise RuntimeError("connection reset in trailing frame")
|
||||
|
||||
def test_trailing_frame_failure_skips_fallback(self, patch_model_utils):
|
||||
backup = FakeLLM(stream_chunks=["fallback chunk"])
|
||||
patch_model_utils(
|
||||
get_provider=lambda mid, **_kw: "openai",
|
||||
get_api_key=lambda p: "k",
|
||||
create_llm=lambda type, **kw: backup,
|
||||
)
|
||||
primary = self.FinishThenFailLLM(backup_models=["backup-model"])
|
||||
|
||||
received = []
|
||||
with pytest.raises(RuntimeError, match="trailing frame"):
|
||||
for chunk in primary.gen_stream(**CALL_ARGS):
|
||||
received.append(chunk)
|
||||
|
||||
assert received == ["the full answer"]
|
||||
assert backup.gen_stream_called is False
|
||||
|
||||
def test_pre_finish_failure_still_falls_back(self, patch_model_utils):
|
||||
backup = FakeLLM(stream_chunks=["fallback chunk"])
|
||||
patch_model_utils(
|
||||
get_provider=lambda mid, **_kw: "openai",
|
||||
get_api_key=lambda p: "k",
|
||||
create_llm=lambda type, **kw: backup,
|
||||
)
|
||||
primary = FakeLLM(fail_at=0, backup_models=["backup-model"])
|
||||
out = list(primary.gen_stream(**CALL_ARGS))
|
||||
assert out == ["fallback chunk"]
|
||||
assert backup.gen_stream_called is True
|
||||
Reference in new issue
Block a user