diff --git a/CLAUDE.md b/CLAUDE.md index 5e1a99d..1f5140c 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -58,7 +58,7 @@ Async SQLAlchemy + aiosqlite. `MemoryStore`: schema init (+ lightweight SQLite ` - **`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`), **usage** (`TokenUsage`). -- Usage: API `usage` from completions (stream with `include_usage`); **char≈token fallback** (~4 chars/token) when omitted; `ChatAgent.session_usage` / `last_turn_usage`; CLI end-of-turn + `/status` (marks `(est)` / `(api+est)`). +- Usage: API `usage` from completions (stream with `include_usage`); **char≈token fallback** (~4 chars/token) when omitted; `last_request_usage` (one model call), `last_turn_usage` / `session_usage` (**billed sums** — tool rounds re-send history); CLI end-of-turn + `/status`. - Config ``[agent]``: `system_prompt`, `max_tool_result_chars`, `parallel_tools`, `confirm_destructive`, `path_denylist`, `max_context_tokens` (default 200k est. tokens). ### Tools (`tools/`) diff --git a/src/plyngent/agent/chat.py b/src/plyngent/agent/chat.py index 972c099..1831ff8 100644 --- a/src/plyngent/agent/chat.py +++ b/src/plyngent/agent/chat.py @@ -42,6 +42,8 @@ class ChatAgent: pending_retry_text: str | None session_usage: TokenUsage last_turn_usage: TokenUsage + last_request_usage: TokenUsage + last_turn_rounds: int def __init__( self, @@ -78,6 +80,8 @@ class ChatAgent: self.pending_retry_text = None self.session_usage = TokenUsage() self.last_turn_usage = TokenUsage() + self.last_request_usage = TokenUsage() + self.last_turn_rounds = 0 self._ensure_system_prompt() def _ensure_system_prompt(self) -> None: @@ -123,6 +127,8 @@ class ChatAgent: completed = False turn_usage = TokenUsage() + turn_rounds = 0 + last_request = TokenUsage() try: async for event in run_chat_loop( self.client, @@ -138,6 +144,9 @@ class ChatAgent: max_context_tokens=self.max_context_tokens, ): if isinstance(event, UsageEvent): + # Each tool-loop round re-sends history; sum is billing, not context size. + turn_rounds += 1 + last_request = event.usage turn_usage = turn_usage.add(event.usage) self.session_usage = self.session_usage.add(event.usage) yield event @@ -148,6 +157,8 @@ class ChatAgent: raise self.last_turn_usage = turn_usage + self.last_request_usage = last_request + self.last_turn_rounds = turn_rounds for message in self.messages[pre_len:]: await self._persist(message) self.pending_retry_text = None diff --git a/src/plyngent/agent/usage.py b/src/plyngent/agent/usage.py index ab9913b..de6be16 100644 --- a/src/plyngent/agent/usage.py +++ b/src/plyngent/agent/usage.py @@ -40,15 +40,21 @@ class TokenUsage(Struct, omit_defaults=True): def is_zero(self) -> bool: return self.prompt_tokens == 0 and self.completion_tokens == 0 and self.total_tokens == 0 - def format_line(self) -> str: + def format_line(self, *, billed: bool = False) -> str: + """Format for display. + + When ``billed`` is True, label as cumulative API billing (sum of rounds), + not a single context snapshot — each tool-loop round re-sends history. + """ tag = "" if self.source == "estimate": tag = " (est)" elif self.source == "mixed": tag = " (api+est)" + prefix = "billed " if billed else "" return ( - f"tokens prompt={self.prompt_tokens} completion={self.completion_tokens} " - f"total={self.total_tokens}{tag}" + f"{prefix}tokens prompt={self.prompt_tokens} " + f"completion={self.completion_tokens} total={self.total_tokens}{tag}" ) diff --git a/src/plyngent/cli/repl.py b/src/plyngent/cli/repl.py index 0263ccc..dc98ebf 100644 --- a/src/plyngent/cli/repl.py +++ b/src/plyngent/cli/repl.py @@ -63,17 +63,22 @@ def _cmd_status(state: ReplState) -> None: ctx_budget = state.agent.max_context_tokens session_u = state.agent.session_usage last_u = state.agent.last_turn_usage + last_req = state.agent.last_request_usage + last_rounds = state.agent.last_turn_rounds click.echo( f"provider={state.provider_name} model={state.model}\n" f"session={state.session_id} messages={len(state.agent.messages)} " f"pending_retry={pending_disp}\n" f"tools={'on' if state.tools_enabled else 'off'} " f"rounds={state.max_rounds} stream={'on' if state.agent.stream else 'off'}\n" - f"context_tokens~={ctx_tokens}/{ctx_budget} (est) " + f"context_tokens~={ctx_tokens}/{ctx_budget} (est, once) " f"context_chars={ctx_chars} " 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"last_request={last_req.format_line()} " + f"(last model call; ~context size if from API)\n" + f"usage_last_turn={last_u.format_line(billed=True)} " + f"rounds={last_rounds}\n" + f"usage_session={session_u.format_line(billed=True)}\n" f"workspace={state.workspace}" ) diff --git a/src/plyngent/cli/retry.py b/src/plyngent/cli/retry.py index cad2067..256de54 100644 --- a/src/plyngent/cli/retry.py +++ b/src/plyngent/cli/retry.py @@ -89,8 +89,19 @@ async def run_cancellable[T](coro: Coroutine[object, object, T]) -> T: def _echo_turn_usage(agent: ChatAgent) -> None: + if agent.last_turn_usage.is_zero() and agent.last_request_usage.is_zero(): + return + rounds = agent.last_turn_rounds + parts: list[str] = [] + if not agent.last_request_usage.is_zero(): + parts.append(f"last_request {agent.last_request_usage.format_line()}") if not agent.last_turn_usage.is_zero(): - click.secho(f"[{agent.last_turn_usage.format_line()}]", fg="bright_black") + label = agent.last_turn_usage.format_line(billed=True) + if rounds > 1: + parts.append(f"turn {label} over {rounds} rounds") + else: + parts.append(f"turn {label}") + click.secho(f"[{'; '.join(parts)}]", fg="bright_black") async def _wait_for_retry(attempt: int, max_retries: int, wait: float) -> bool: diff --git a/tests/test_agent/test_loop.py b/tests/test_agent/test_loop.py index 1f2d967..23709ee 100644 --- a/tests/test_agent/test_loop.py +++ b/tests/test_agent/test_loop.py @@ -230,12 +230,51 @@ async def test_chat_agent_accumulates_session_usage() -> None: 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.last_request_usage.total_tokens == 6 + assert agent.last_turn_rounds == 1 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.last_request_usage.total_tokens == 10 assert agent.session_usage.total_tokens == 16 +async def test_chat_agent_turn_usage_sums_tool_rounds() -> None: + """Multi-round tool loop: turn usage is billing sum; last_request is final call.""" + + @tool + def ping() -> str: + return "pong" + + registry = ToolRegistry([ping]) + client = ScriptedClient( + [ + _response( + AssistantChatMessage( + content="", + tool_calls=[ + AssistantFunctionToolCall( + id="1", + function=AssistantFunctionTool(name="ping", arguments="{}"), + ) + ], + ), + usage={"prompt_tokens": 100, "completion_tokens": 5, "total_tokens": 105}, + ), + _response( + AssistantChatMessage(content="done"), + usage={"prompt_tokens": 200, "completion_tokens": 3, "total_tokens": 203}, + ), + ] + ) + agent = ChatAgent(client, model="m", tools=registry, stream=False) + _ = [e async for e in agent.run("go")] + assert agent.last_turn_rounds == 2 + assert agent.last_request_usage.prompt_tokens == 200 + assert agent.last_turn_usage.prompt_tokens == 300 + assert agent.last_turn_usage.total_tokens == 308 + + async def test_stream_yields_deltas_incrementally() -> None: """Text deltas are yielded as chunks arrive, not only after the full stream.""" diff --git a/tests/test_agent/test_usage.py b/tests/test_agent/test_usage.py index a1bb34f..c2e276b 100644 --- a/tests/test_agent/test_usage.py +++ b/tests/test_agent/test_usage.py @@ -97,3 +97,5 @@ def test_format_line_marks_estimate() -> None: assert "(est)" in line mixed = TokenUsage(prompt_tokens=1, completion_tokens=0, total_tokens=1, source="mixed").format_line() assert "(api+est)" in mixed + billed = TokenUsage(prompt_tokens=1, completion_tokens=0, total_tokens=1).format_line(billed=True) + assert billed.startswith("billed tokens")