From d81e168f3b5652e71d1185996473dcb02734775e Mon Sep 17 00:00:00 2001 From: worldmozara Date: Wed, 15 Jul 2026 10:48:07 +0800 Subject: [PATCH] core/agent+cli: accumulate usage and show on turn end /status ChatAgent tracks session_usage and last_turn_usage; CLI prints a quiet token line after successful turns and includes totals in /status. --- CLAUDE.md | 9 ++++---- src/plyngent/agent/__init__.py | 3 +++ src/plyngent/agent/chat.py | 11 ++++++++++ src/plyngent/cli/display.py | 4 ++++ src/plyngent/cli/repl.py | 4 ++++ src/plyngent/cli/retry.py | 39 ++++++++++++++++++++++------------ 6 files changed, 52 insertions(+), 18 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index c13608d..e254ef5 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -57,7 +57,8 @@ Async SQLAlchemy + aiosqlite. `MemoryStore`: schema init (+ lightweight SQLite ` - **`run_chat_loop`**: multi-round tool loop; default **streaming** text deltas + stream tool-call merge; parallel tools; tool-result char budget; soft context compact on request; cooperative cancel points; optional `on_limit`. - **`ChatAgent`**: optional `MemoryStore` (persist on success only); `stream`; system prompt; `pending_retry_text` + `retry()`. - **`/compact`**: soft-compact tool dumps → model summary (no tools) → **new** session seeded with summary message. -- Events: text_delta, assistant_message, tool_call/result, max_rounds, **error** (`retryable`/`source`), **cancelled** (`reason`). +- Events: text_delta, assistant_message, tool_call/result, max_rounds, **error** (`retryable`/`source`), **cancelled** (`reason`), **usage** (`TokenUsage`). +- Usage: API `usage` from completions (stream with `include_usage`); `ChatAgent.session_usage` / `last_turn_usage`; CLI end-of-turn + `/status`. - Config ``[agent]``: `system_prompt`, `max_tool_result_chars`, `parallel_tools`, `confirm_destructive`, `path_denylist`, `max_context_chars`. ### Tools (`tools/`) @@ -97,10 +98,10 @@ Basedpyright `recommended`. Ruff includes `ANN` (private return types `ANN202` i ## Roadmap notes (single-user → platform) -- **Phase D (context quality)**: soft char budget, request compact, `/compact`, richer errors/cancel, workspace sessions — **current**. Context size is **char estimate only** until usage lands. -- **No local tokenizer stage** for now. Prefer **API usage v2** later. +- **Phase D (context quality)**: soft char budget, request compact, `/compact`, richer errors/cancel, workspace sessions. Context size is **char estimate** plus optional API usage when reported. +- **No local tokenizer stage** for now. - **Phase E**: tooling depth (grep/glob, VCS backends; prefer `edit_replace` / `edit_lineno` over model-generated patches). -- **Phase F (providers + usage v2)**: capture response/stream `usage` (prompt/completion/total); session/turn totals; `/status` or end-of-turn line; optional cost. Optional later: tokenizer-backed estimates if needed. +- **Phase F (providers + usage v2)**: capture response/stream `usage` — **in progress / landing**; session/turn totals; `/status` + end-of-turn line. Optional later: cost, tokenizer-backed estimates. - **Phase G–H**: CLI polish, hardening; then multi-tenant platform (`router/`, auth, sandboxed tools). ## Commit messages diff --git a/src/plyngent/agent/__init__.py b/src/plyngent/agent/__init__.py index a0ff27d..8f46788 100644 --- a/src/plyngent/agent/__init__.py +++ b/src/plyngent/agent/__init__.py @@ -10,6 +10,7 @@ from .events import MaxRoundsEvent as MaxRoundsEvent from .events import TextDeltaEvent as TextDeltaEvent from .events import ToolCallEvent as ToolCallEvent from .events import ToolResultEvent as ToolResultEvent +from .events import UsageEvent as UsageEvent from .loop import DEFAULT_MAX_ROUNDS as DEFAULT_MAX_ROUNDS from .loop import run_chat_loop as run_chat_loop from .tools import DangerClassifier as DangerClassifier @@ -18,3 +19,5 @@ from .tools import ToolDefinition as ToolDefinition 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 token_usage_from_api as token_usage_from_api diff --git a/src/plyngent/agent/chat.py b/src/plyngent/agent/chat.py index 648186e..ec79948 100644 --- a/src/plyngent/agent/chat.py +++ b/src/plyngent/agent/chat.py @@ -5,7 +5,9 @@ from typing import TYPE_CHECKING from plyngent.lmproto.openai_compatible.model import SystemChatMessage, UserChatMessage from .budget import DEFAULT_CONTEXT_MAX_CHARS, DEFAULT_TOOL_RESULT_MAX_CHARS +from .events import UsageEvent from .loop import DEFAULT_MAX_ROUNDS, run_chat_loop +from .usage import TokenUsage if TYPE_CHECKING: from collections.abc import AsyncIterator, Awaitable, Callable, Sequence @@ -38,6 +40,8 @@ class ChatAgent: max_context_chars: int messages: list[AnyChatMessage] pending_retry_text: str | None + session_usage: TokenUsage + last_turn_usage: TokenUsage def __init__( self, @@ -72,6 +76,8 @@ class ChatAgent: self.max_context_chars = max_context_chars self.messages = list(messages) if messages is not None else [] self.pending_retry_text = None + self.session_usage = TokenUsage() + self.last_turn_usage = TokenUsage() self._ensure_system_prompt() def _ensure_system_prompt(self) -> None: @@ -116,6 +122,7 @@ class ChatAgent: self.messages.append(user_msg) completed = False + turn_usage = TokenUsage() try: async for event in run_chat_loop( self.client, @@ -130,6 +137,9 @@ class ChatAgent: parallel_tools=self.parallel_tools, max_context_chars=self.max_context_chars, ): + if isinstance(event, UsageEvent): + turn_usage = turn_usage.add(event.usage) + self.session_usage = self.session_usage.add(event.usage) yield event completed = True except BaseException: @@ -137,6 +147,7 @@ class ChatAgent: self._rollback_turn(pre_len, user_msg.content) raise + self.last_turn_usage = turn_usage for message in self.messages[pre_len:]: await self._persist(message) self.pending_retry_text = None diff --git a/src/plyngent/cli/display.py b/src/plyngent/cli/display.py index 196c966..9c087e9 100644 --- a/src/plyngent/cli/display.py +++ b/src/plyngent/cli/display.py @@ -11,6 +11,7 @@ from plyngent.agent import ( TextDeltaEvent, ToolCallEvent, ToolResultEvent, + UsageEvent, ) from plyngent.lmproto.openai_compatible.model import AssistantFunctionToolCall @@ -82,6 +83,9 @@ async def render_events(events: AsyncIterator[AgentEvent]) -> None: # noqa: C90 ) else: click.secho(f"\n[max rounds reached: {event.rounds}]", fg="red") + elif isinstance(event, UsageEvent): + # Printed at end of turn as a summary; individual round usage is quiet. + _ = event else: # AssistantMessageEvent — text already shown via TextDeltaEvent. _ = event diff --git a/src/plyngent/cli/repl.py b/src/plyngent/cli/repl.py index 13a7fdb..1c536a4 100644 --- a/src/plyngent/cli/repl.py +++ b/src/plyngent/cli/repl.py @@ -60,6 +60,8 @@ def _cmd_status(state: ReplState) -> None: pending_disp = "yes" if pending else "no" ctx_chars = estimate_messages_chars(state.agent.messages) ctx_budget = state.agent.max_context_chars + session_u = state.agent.session_usage + last_u = state.agent.last_turn_usage click.echo( f"provider={state.provider_name} model={state.model}\n" f"session={state.session_id} messages={len(state.agent.messages)} " @@ -68,6 +70,8 @@ def _cmd_status(state: ReplState) -> None: f"rounds={state.max_rounds} stream={'on' if state.agent.stream else 'off'}\n" f"context_chars={ctx_chars}/{ctx_budget} " f"tool_result_max={state.agent.max_tool_result_chars}\n" + f"usage_session={session_u.format_line()}\n" + f"usage_last_turn={last_u.format_line()}\n" f"workspace={state.workspace}" ) diff --git a/src/plyngent/cli/retry.py b/src/plyngent/cli/retry.py index 2efe61f..cad2067 100644 --- a/src/plyngent/cli/retry.py +++ b/src/plyngent/cli/retry.py @@ -88,6 +88,29 @@ async def run_cancellable[T](coro: Coroutine[object, object, T]) -> T: await task +def _echo_turn_usage(agent: ChatAgent) -> None: + if not agent.last_turn_usage.is_zero(): + click.secho(f"[{agent.last_turn_usage.format_line()}]", fg="bright_black") + + +async def _wait_for_retry(attempt: int, max_retries: int, wait: float) -> bool: + click.secho( + f"auto-retry {attempt}/{max_retries} in {wait:g}s " + f"(Ctrl+C to cancel; then /retry later)", + fg="yellow", + ) + try: + ok = await sleep_cancellable(wait) + except KeyboardInterrupt: + ok = False + if not ok: + click.secho("auto-retry cancelled; use /retry to try again", fg="yellow") + click.echo() + return False + click.secho(f"retrying ({attempt}/{max_retries})…", fg="yellow") + return True + + async def run_turn_with_retries( agent: ChatAgent, *, @@ -131,24 +154,12 @@ async def run_turn_with_retries( ) click.echo() return False - wait = delays[attempt] attempt += 1 - click.secho( - f"auto-retry {attempt}/{max_retries} in {wait:g}s " - f"(Ctrl+C to cancel; then /retry later)", - fg="yellow", - ) - try: - ok = await sleep_cancellable(wait) - except KeyboardInterrupt: - ok = False - if not ok: - click.secho("auto-retry cancelled; use /retry to try again", fg="yellow") - click.echo() + if not await _wait_for_retry(attempt, max_retries, wait): return False - click.secho(f"retrying ({attempt}/{max_retries})…", fg="yellow") else: + _echo_turn_usage(agent) return True