fix usage caturing

This commit is contained in:
9qeklajc
2026-06-03 17:04:48 +02:00
parent 320efe2e85
commit d4296d6087
2 changed files with 72 additions and 0 deletions
+27
View File
@@ -790,6 +790,33 @@ class BaseUpstreamProvider:
self._current_stream_id = f"chatcmpl-{uuid.uuid4()}"
obj["id"] = self._current_stream_id
if isinstance(obj.get("usage"), dict):
# Capture usage for end-of-stream cost reconciliation.
# Some models (e.g. Gemini thinking models over the
# OpenAI-compat endpoint) attach ``usage`` to the SAME
# chunk that carries the final content/finish_reason
# rather than sending a separate ``choices: []`` usage
# chunk. Only swallow the chunk when it is a pure usage
# chunk (no choices); otherwise the content would be
# silently dropped and the client would receive no
# assistant message at all.
if obj.get("choices"):
# Capture usage (with model) for the cost trailer,
# but with choices stripped so the trailer never
# re-emits this chunk's content.
usage_chunk_data = {
k: v for k, v in obj.items() if k != "choices"
}
usage_chunk_data["choices"] = []
# Forward the content now, without usage, so token
# usage is reported exactly once (in the trailer).
forward = {k: v for k, v in obj.items() if k != "usage"}
yield (
prefix
+ b"data: "
+ json.dumps(forward).encode()
+ b"\n\n"
)
return
usage_chunk_data = obj
return
yield prefix + b"data: " + json.dumps(obj).encode() + b"\n\n"
@@ -254,6 +254,51 @@ async def test_openrouter_mid_stream_error_event() -> None:
assert any("error" in o for o in objs), "error event must be forwarded intact"
@pytest.mark.asyncio
async def test_gemini_combined_content_and_usage_chunk() -> None:
"""Gemini thinking models pack usage into the final *content* chunk.
Regression: the parser swallowed any chunk carrying a ``usage`` dict, so
when content + usage arrived together the assistant text was dropped and
the client saw "no assistant messages" despite a 200 + token accounting.
"""
chunks = [
b'data: {"id":"g","choices":[{"delta":{"content":"the answer"},'
b'"finish_reason":"stop"}],"usage":{"prompt_tokens":3,'
b'"completion_tokens":2,"total_tokens":5}}\n\n',
b"data: [DONE]\n\n",
]
out = await _drive(chunks)
objs = _assert_clean(out)
contents = [
c["delta"]["content"]
for o in objs
for c in o.get("choices", [])
if "delta" in c
]
# Content delivered exactly once (not dropped, not duplicated by the trailer).
assert contents == ["the answer"]
@pytest.mark.asyncio
async def test_separate_usage_chunk_not_forwarded_as_content() -> None:
"""A pure usage chunk (choices: []) is still swallowed, content intact."""
chunks = [
b'data: {"id":"x","choices":[{"delta":{"content":"hello"}}]}\n\n',
b'data: {"id":"x","choices":[],"usage":{"total_tokens":4}}\n\n',
b"data: [DONE]\n\n",
]
out = await _drive(chunks)
objs = _assert_clean(out)
contents = [
c["delta"]["content"]
for o in objs
for c in o.get("choices", [])
if "delta" in c
]
assert contents == ["hello"]
@pytest.mark.asyncio
async def test_requested_model_override_applied() -> None:
"""Model rewriting still works through the buffered parser."""