From 28b6767b32cc2644dae5e63b66d805492089efc7 Mon Sep 17 00:00:00 2001 From: worldmozara Date: Wed, 15 Jul 2026 10:47:59 +0800 Subject: [PATCH] core/agent: emit UsageEvent from stream and non-stream rounds Request stream_options.include_usage; yield usage after each model completion when the provider reports it. --- src/plyngent/agent/events.py | 9 +++++++ src/plyngent/agent/loop.py | 36 ++++++++++++++++++------- tests/test_agent/test_loop.py | 49 +++++++++++++++++++++++++++++++++-- 3 files changed, 82 insertions(+), 12 deletions(-) diff --git a/src/plyngent/agent/events.py b/src/plyngent/agent/events.py index 0825379..3ff38a7 100644 --- a/src/plyngent/agent/events.py +++ b/src/plyngent/agent/events.py @@ -6,6 +6,8 @@ from plyngent.lmproto.openai_compatible.model import ( # noqa: TC001 ToolChatMessage, ) +from .usage import TokenUsage # noqa: TC001 + class TextDeltaEvent(Struct, tag_field="type", tag="text_delta"): content: str @@ -38,6 +40,12 @@ class CancelledEvent(Struct, tag_field="type", tag="cancelled"): reason: str = "" +class UsageEvent(Struct, tag_field="type", tag="usage"): + """Token usage for one model completion (one tool-loop round).""" + + usage: TokenUsage + + type AgentEvent = ( TextDeltaEvent | AssistantMessageEvent @@ -46,4 +54,5 @@ type AgentEvent = ( | MaxRoundsEvent | ErrorEvent | CancelledEvent + | UsageEvent ) diff --git a/src/plyngent/agent/loop.py b/src/plyngent/agent/loop.py index 9df8e99..602e1f2 100644 --- a/src/plyngent/agent/loop.py +++ b/src/plyngent/agent/loop.py @@ -4,6 +4,7 @@ import asyncio import inspect from typing import TYPE_CHECKING, cast +import msgspec from msgspec import UNSET from plyngent.lmproto.openai_compatible.client import merge_stream_tool_calls @@ -12,6 +13,7 @@ from plyngent.lmproto.openai_compatible.model import ( AssistantChatMessage, AssistantFunctionToolCall, ChatCompletionsParam, + StreamOptions, StreamToolCallDelta, ToolChatMessage, ) @@ -31,7 +33,9 @@ from .events import ( TextDeltaEvent, ToolCallEvent, ToolResultEvent, + UsageEvent, ) +from .usage import token_usage_from_api if TYPE_CHECKING: from collections.abc import AsyncIterator, Awaitable, Callable, Sequence @@ -110,19 +114,21 @@ async def _execute_tool_calls( yield ToolResultEvent(message=tool_msg) -async def _non_stream_assistant( +async def _non_stream_round( client: ChatClient, param: ChatCompletionsParam, -) -> tuple[AssistantChatMessage, list[TextDeltaEvent]]: +) -> AsyncIterator[AgentEvent]: response = await client.chat_completions(param, stream=False) if not response.choices: msg = "chat completion response contained no choices" raise RuntimeError(msg) assistant = response.choices[0].message - text_events: list[TextDeltaEvent] = [] if isinstance(assistant.content, str) and assistant.content: - text_events.append(TextDeltaEvent(content=assistant.content)) - return assistant, text_events + yield TextDeltaEvent(content=assistant.content) + yield AssistantMessageEvent(message=assistant) + usage = token_usage_from_api(response.usage) + if usage is not None: + yield UsageEvent(usage=usage) async def _stream_round( @@ -133,12 +139,21 @@ async def _stream_round( Pattern: ``stream = await client.chat_completions(..., stream=True)`` then ``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. """ - stream = await client.chat_completions(param, stream=True) + stream_param = msgspec.structs.replace( + param, + stream_options=StreamOptions(include_usage=True), + ) + stream = await client.chat_completions(stream_param, stream=True) content_parts: list[str] = [] tool_deltas: list[StreamToolCallDelta] = [] + last_usage = None async for chunk in stream: + usage = token_usage_from_api(chunk.usage) + if usage is not None: + last_usage = usage if not chunk.choices: continue choice = chunk.choices[0] @@ -161,6 +176,8 @@ async def _stream_round( tool_calls=tool_calls, ) yield AssistantMessageEvent(message=assistant) + if last_usage is not None: + yield UsageEvent(usage=last_usage) async def _assistant_round( @@ -177,11 +194,10 @@ async def _assistant_round( messages.append(event.message) yield event return - assistant, text_events = await _non_stream_assistant(client, param) - for event in text_events: + async for event in _non_stream_round(client, param): + if isinstance(event, AssistantMessageEvent): + messages.append(event.message) yield event - messages.append(assistant) - yield AssistantMessageEvent(message=assistant) def _last_assistant(messages: list[AnyChatMessage], pre_len: int) -> AssistantChatMessage: diff --git a/tests/test_agent/test_loop.py b/tests/test_agent/test_loop.py index c2e5290..14f61a9 100644 --- a/tests/test_agent/test_loop.py +++ b/tests/test_agent/test_loop.py @@ -13,6 +13,7 @@ from plyngent.agent import ( ToolCallEvent, ToolRegistry, ToolResultEvent, + UsageEvent, run_chat_loop, tool, ) @@ -146,7 +147,11 @@ class ScriptedClient: yield chunk -def _response(message: AssistantChatMessage) -> ChatCompletionResponse: +def _response( + message: AssistantChatMessage, + *, + usage: dict[str, int] | None = None, +) -> ChatCompletionResponse: return ChatCompletionResponse( id="1", object="chat.completion", @@ -161,7 +166,7 @@ def _response(message: AssistantChatMessage) -> ChatCompletionResponse: ) ], system_fingerprint="", - usage={}, + usage=usage if usage is not None else {}, ) @@ -180,6 +185,46 @@ async def test_run_chat_loop_text_only() -> None: assert len(client.calls) == 1 +async def test_non_stream_emits_usage() -> None: + client = ScriptedClient( + [ + _response( + AssistantChatMessage(content="hi"), + usage={"prompt_tokens": 9, "completion_tokens": 2, "total_tokens": 11}, + ), + ] + ) + messages: list[AnyChatMessage] = [UserChatMessage(content="x")] + 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.prompt_tokens == 9 + assert usages[0].usage.completion_tokens == 2 + assert usages[0].usage.total_tokens == 11 + + +async def test_chat_agent_accumulates_session_usage() -> None: + client = ScriptedClient( + [ + _response( + AssistantChatMessage(content="a"), + usage={"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6}, + ), + _response( + AssistantChatMessage(content="b"), + usage={"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10}, + ), + ] + ) + agent = ChatAgent(client, model="m", stream=False) + _ = [e async for e in agent.run("one")] + assert agent.last_turn_usage.total_tokens == 6 + assert agent.session_usage.total_tokens == 6 + _ = [e async for e in agent.run("two")] + assert agent.last_turn_usage.total_tokens == 10 + assert agent.session_usage.total_tokens == 16 + + async def test_stream_yields_deltas_incrementally() -> None: """Text deltas are yielded as chunks arrive, not only after the full stream."""