mirror of
https://github.com/NCBM/plyngent.git
synced 2026-07-23 22:44:59 +08:00
core/agent: always emit UsageEvent with api or estimate
Every model round yields usage; prefer provider usage, else estimate from request messages and assistant content.
This commit is contained in:
@@ -20,4 +20,6 @@ from .tools import ToolRegistry as ToolRegistry
|
|||||||
from .tools import schema_from_callable as schema_from_callable
|
from .tools import schema_from_callable as schema_from_callable
|
||||||
from .tools import tool as tool
|
from .tools import tool as tool
|
||||||
from .usage import TokenUsage as TokenUsage
|
from .usage import TokenUsage as TokenUsage
|
||||||
|
from .usage import chars_to_tokens as chars_to_tokens
|
||||||
|
from .usage import estimate_token_usage as estimate_token_usage
|
||||||
from .usage import token_usage_from_api as token_usage_from_api
|
from .usage import token_usage_from_api as token_usage_from_api
|
||||||
|
|||||||
@@ -41,7 +41,10 @@ class CancelledEvent(Struct, tag_field="type", tag="cancelled"):
|
|||||||
|
|
||||||
|
|
||||||
class UsageEvent(Struct, tag_field="type", tag="usage"):
|
class UsageEvent(Struct, tag_field="type", tag="usage"):
|
||||||
"""Token usage for one model completion (one tool-loop round)."""
|
"""Token usage for one model completion (one tool-loop round).
|
||||||
|
|
||||||
|
``source`` is mirrored from ``usage.source`` (``api`` / ``estimate``).
|
||||||
|
"""
|
||||||
|
|
||||||
usage: TokenUsage
|
usage: TokenUsage
|
||||||
|
|
||||||
|
|||||||
+12
-10
@@ -35,7 +35,7 @@ from .events import (
|
|||||||
ToolResultEvent,
|
ToolResultEvent,
|
||||||
UsageEvent,
|
UsageEvent,
|
||||||
)
|
)
|
||||||
from .usage import token_usage_from_api
|
from .usage import resolve_round_usage, token_usage_from_api
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from collections.abc import AsyncIterator, Awaitable, Callable, Sequence
|
from collections.abc import AsyncIterator, Awaitable, Callable, Sequence
|
||||||
@@ -126,9 +126,9 @@ async def _non_stream_round(
|
|||||||
if isinstance(assistant.content, str) and assistant.content:
|
if isinstance(assistant.content, str) and assistant.content:
|
||||||
yield TextDeltaEvent(content=assistant.content)
|
yield TextDeltaEvent(content=assistant.content)
|
||||||
yield AssistantMessageEvent(message=assistant)
|
yield AssistantMessageEvent(message=assistant)
|
||||||
usage = token_usage_from_api(response.usage)
|
yield UsageEvent(
|
||||||
if usage is not None:
|
usage=resolve_round_usage(response.usage, param.messages, assistant),
|
||||||
yield UsageEvent(usage=usage)
|
)
|
||||||
|
|
||||||
|
|
||||||
async def _stream_round(
|
async def _stream_round(
|
||||||
@@ -140,6 +140,7 @@ async def _stream_round(
|
|||||||
Pattern: ``stream = await client.chat_completions(..., stream=True)`` then
|
Pattern: ``stream = await client.chat_completions(..., stream=True)`` then
|
||||||
``async for chunk in stream``. Tool-call deltas are merged after the stream.
|
``async for chunk in stream``. Tool-call deltas are merged after the stream.
|
||||||
Requests ``stream_options.include_usage`` so providers may send a final usage chunk.
|
Requests ``stream_options.include_usage`` so providers may send a final usage chunk.
|
||||||
|
Falls back to a char-based token estimate when the provider omits usage.
|
||||||
"""
|
"""
|
||||||
stream_param = msgspec.structs.replace(
|
stream_param = msgspec.structs.replace(
|
||||||
param,
|
param,
|
||||||
@@ -148,12 +149,12 @@ async def _stream_round(
|
|||||||
stream = await client.chat_completions(stream_param, stream=True)
|
stream = await client.chat_completions(stream_param, stream=True)
|
||||||
content_parts: list[str] = []
|
content_parts: list[str] = []
|
||||||
tool_deltas: list[StreamToolCallDelta] = []
|
tool_deltas: list[StreamToolCallDelta] = []
|
||||||
last_usage = None
|
last_api_usage: object = UNSET
|
||||||
|
|
||||||
async for chunk in stream:
|
async for chunk in stream:
|
||||||
usage = token_usage_from_api(chunk.usage)
|
parsed = token_usage_from_api(chunk.usage)
|
||||||
if usage is not None:
|
if parsed is not None:
|
||||||
last_usage = usage
|
last_api_usage = chunk.usage
|
||||||
if not chunk.choices:
|
if not chunk.choices:
|
||||||
continue
|
continue
|
||||||
choice = chunk.choices[0]
|
choice = chunk.choices[0]
|
||||||
@@ -176,8 +177,9 @@ async def _stream_round(
|
|||||||
tool_calls=tool_calls,
|
tool_calls=tool_calls,
|
||||||
)
|
)
|
||||||
yield AssistantMessageEvent(message=assistant)
|
yield AssistantMessageEvent(message=assistant)
|
||||||
if last_usage is not None:
|
yield UsageEvent(
|
||||||
yield UsageEvent(usage=last_usage)
|
usage=resolve_round_usage(last_api_usage, param.messages, assistant),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
async def _assistant_round(
|
async def _assistant_round(
|
||||||
|
|||||||
@@ -201,6 +201,17 @@ async def test_non_stream_emits_usage() -> None:
|
|||||||
assert usages[0].usage.prompt_tokens == 9
|
assert usages[0].usage.prompt_tokens == 9
|
||||||
assert usages[0].usage.completion_tokens == 2
|
assert usages[0].usage.completion_tokens == 2
|
||||||
assert usages[0].usage.total_tokens == 11
|
assert usages[0].usage.total_tokens == 11
|
||||||
|
assert usages[0].usage.source == "api"
|
||||||
|
|
||||||
|
|
||||||
|
async def test_non_stream_estimates_usage_when_missing() -> None:
|
||||||
|
client = ScriptedClient([_response(AssistantChatMessage(content="hello"))])
|
||||||
|
messages: list[AnyChatMessage] = [UserChatMessage(content="hi")]
|
||||||
|
events = [e async for e in run_chat_loop(client, messages, model="m", stream=False)]
|
||||||
|
usages = [e for e in events if isinstance(e, UsageEvent)]
|
||||||
|
assert len(usages) == 1
|
||||||
|
assert usages[0].usage.source == "estimate"
|
||||||
|
assert usages[0].usage.total_tokens > 0
|
||||||
|
|
||||||
|
|
||||||
async def test_chat_agent_accumulates_session_usage() -> None:
|
async def test_chat_agent_accumulates_session_usage() -> None:
|
||||||
|
|||||||
Reference in New Issue
Block a user