mirror of
https://github.com/NCBM/plyngent.git
synced 2026-11-07 23:47:54 +08:00
core/agent: count the prompt a search turn really sent
The compat opens with the question alone (``message_start`` reports 165 tokens) and puts what the search then pulled in on ``message_delta`` (12468); merging preferred the earlier value, so the CLI's context gauge and the auto-compact threshold ran ~12k short of the request. The larger of the two wins now (output tokens are the delta's cumulative count either way).
This commit is contained in:
@@ -66,11 +66,17 @@ class _ToolBlockState:
|
|||||||
|
|
||||||
|
|
||||||
def _merge_usage(previous: dict[str, Any] | None, new: dict[str, Any]) -> dict[str, Any]:
|
def _merge_usage(previous: dict[str, Any] | None, new: dict[str, Any]) -> dict[str, Any]:
|
||||||
"""Merge message_start input tokens with message_delta output tokens."""
|
"""Merge ``message_start`` usage with the later ``message_delta`` usage.
|
||||||
|
|
||||||
|
DeepSeek's compat reports a search turn's real prompt tokens only on the
|
||||||
|
delta (``message_start`` sees the question alone, the delta counts what the
|
||||||
|
search pulled in), so the larger of the two wins; output tokens are the
|
||||||
|
delta's cumulative count either way.
|
||||||
|
"""
|
||||||
if previous is None:
|
if previous is None:
|
||||||
return new
|
return new
|
||||||
prompt = int(previous.get("prompt_tokens", 0) or new.get("prompt_tokens", 0) or 0)
|
prompt = max(int(previous.get("prompt_tokens", 0) or 0), int(new.get("prompt_tokens", 0) or 0))
|
||||||
completion = int(new.get("completion_tokens", 0) or previous.get("completion_tokens", 0) or 0)
|
completion = max(int(previous.get("completion_tokens", 0) or 0), int(new.get("completion_tokens", 0) or 0))
|
||||||
return {
|
return {
|
||||||
"prompt_tokens": prompt,
|
"prompt_tokens": prompt,
|
||||||
"completion_tokens": completion,
|
"completion_tokens": completion,
|
||||||
|
|||||||
@@ -497,3 +497,38 @@ async def test_dispatch_messages_without_provider_tools_sends_no_server_tool() -
|
|||||||
param = ChatCompletionsParam(model="deepseek-flash", messages=[UserChatMessage(content="hi")])
|
param = ChatCompletionsParam(model="deepseek-flash", messages=[UserChatMessage(content="hi")])
|
||||||
_ = await dispatch_messages(cast("Any", client), param, provider_tools=[])
|
_ = await dispatch_messages(cast("Any", client), param, provider_tools=[])
|
||||||
assert client.calls[0].tools is UNSET
|
assert client.calls[0].tools is UNSET
|
||||||
|
|
||||||
|
|
||||||
|
async def test_search_turn_reports_the_prompt_the_delta_counts() -> None:
|
||||||
|
"""A search turn's real prompt size only shows up on ``message_delta``.
|
||||||
|
|
||||||
|
DeepSeek's compat opens with the question alone (``message_start`` sees 165
|
||||||
|
tokens) and reports what the search then pulled in on the delta (12468);
|
||||||
|
keeping the smaller one leaves the CLI's context gauge — and the
|
||||||
|
auto-compact threshold — that far short of what the model was sent.
|
||||||
|
"""
|
||||||
|
events: list[AnthropicStreamEvent] = [
|
||||||
|
AnthropicMessageStart(
|
||||||
|
message=AnthropicMessageResponse(
|
||||||
|
id="msg_s",
|
||||||
|
model="deepseek-flash",
|
||||||
|
content=[],
|
||||||
|
usage=AnthropicUsage(input_tokens=165, output_tokens=0),
|
||||||
|
)
|
||||||
|
),
|
||||||
|
AnthropicContentBlockStart(index=0, content_block=AnthropicRawContentBlock(type="text", text="")),
|
||||||
|
AnthropicContentBlockDelta(index=0, delta=AnthropicRawContentBlock(type="text_delta", text="done")),
|
||||||
|
AnthropicMessageDelta(
|
||||||
|
delta={"stop_reason": "end_turn"},
|
||||||
|
usage=AnthropicUsage(input_tokens=12468, output_tokens=198),
|
||||||
|
),
|
||||||
|
AnthropicMessageStop(),
|
||||||
|
]
|
||||||
|
client = ScriptedMessagesClient(stream_events=[events])
|
||||||
|
param = ChatCompletionsParam(model="deepseek-flash", messages=[UserChatMessage(content="hi")])
|
||||||
|
stream = await dispatch_messages(cast("Any", client), param, stream=True)
|
||||||
|
chunks = [chunk async for chunk in cast("Any", stream)]
|
||||||
|
usage = next(getattr(chunk, "usage", None) for chunk in chunks if getattr(chunk, "usage", None))
|
||||||
|
assert isinstance(usage, dict)
|
||||||
|
assert usage["prompt_tokens"] == 12468
|
||||||
|
assert usage["completion_tokens"] == 198
|
||||||
|
|||||||
Reference in New Issue
Block a user