diff --git a/src/plyngent/agent/__init__.py b/src/plyngent/agent/__init__.py index 8f46788..393f1fe 100644 --- a/src/plyngent/agent/__init__.py +++ b/src/plyngent/agent/__init__.py @@ -20,4 +20,6 @@ from .tools import ToolRegistry as ToolRegistry from .tools import schema_from_callable as schema_from_callable from .tools import tool as tool 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 diff --git a/src/plyngent/agent/events.py b/src/plyngent/agent/events.py index 3ff38a7..7d47eb6 100644 --- a/src/plyngent/agent/events.py +++ b/src/plyngent/agent/events.py @@ -41,7 +41,10 @@ class CancelledEvent(Struct, tag_field="type", tag="cancelled"): 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 diff --git a/src/plyngent/agent/loop.py b/src/plyngent/agent/loop.py index 602e1f2..3ddeaba 100644 --- a/src/plyngent/agent/loop.py +++ b/src/plyngent/agent/loop.py @@ -35,7 +35,7 @@ from .events import ( ToolResultEvent, UsageEvent, ) -from .usage import token_usage_from_api +from .usage import resolve_round_usage, token_usage_from_api if TYPE_CHECKING: 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: 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) + yield UsageEvent( + usage=resolve_round_usage(response.usage, param.messages, assistant), + ) async def _stream_round( @@ -140,6 +140,7 @@ 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. + Falls back to a char-based token estimate when the provider omits usage. """ stream_param = msgspec.structs.replace( param, @@ -148,12 +149,12 @@ async def _stream_round( stream = await client.chat_completions(stream_param, stream=True) content_parts: list[str] = [] tool_deltas: list[StreamToolCallDelta] = [] - last_usage = None + last_api_usage: object = UNSET async for chunk in stream: - usage = token_usage_from_api(chunk.usage) - if usage is not None: - last_usage = usage + parsed = token_usage_from_api(chunk.usage) + if parsed is not None: + last_api_usage = chunk.usage if not chunk.choices: continue choice = chunk.choices[0] @@ -176,8 +177,9 @@ async def _stream_round( tool_calls=tool_calls, ) yield AssistantMessageEvent(message=assistant) - if last_usage is not None: - yield UsageEvent(usage=last_usage) + yield UsageEvent( + usage=resolve_round_usage(last_api_usage, param.messages, assistant), + ) async def _assistant_round( diff --git a/tests/test_agent/test_loop.py b/tests/test_agent/test_loop.py index 14f61a9..1f2d967 100644 --- a/tests/test_agent/test_loop.py +++ b/tests/test_agent/test_loop.py @@ -201,6 +201,17 @@ async def test_non_stream_emits_usage() -> None: assert usages[0].usage.prompt_tokens == 9 assert usages[0].usage.completion_tokens == 2 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: