diff --git a/CLAUDE.md b/CLAUDE.md index 1f5140c..0c5d3b4 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; `last_request_usage` (one model call), `last_turn_usage` / `session_usage` (**billed sums** — tool rounds re-send history); CLI end-of-turn + `/status`. +- Usage: API `usage` from completions (stream with `include_usage`); **char≈token fallback** (~4 chars/token) when omitted; **context size** = last request ``prompt_tokens`` (API preferred); `last_turn_usage` / `session_usage` are **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 1831ff8..23158c3 100644 --- a/src/plyngent/agent/chat.py +++ b/src/plyngent/agent/chat.py @@ -4,7 +4,11 @@ from typing import TYPE_CHECKING from plyngent.lmproto.openai_compatible.model import SystemChatMessage, UserChatMessage -from .budget import DEFAULT_CONTEXT_MAX_TOKENS, DEFAULT_TOOL_RESULT_MAX_CHARS +from .budget import ( + DEFAULT_CONTEXT_MAX_TOKENS, + DEFAULT_TOOL_RESULT_MAX_CHARS, + estimate_messages_tokens, +) from .events import UsageEvent from .loop import DEFAULT_MAX_ROUNDS, run_chat_loop from .usage import TokenUsage @@ -84,6 +88,25 @@ class ChatAgent: self.last_turn_rounds = 0 self._ensure_system_prompt() + @property + def context_tokens(self) -> int: + """Best current context size (tokens). + + Prefers the last model call's ``prompt_tokens`` (API or per-request + estimate) — that is the real size of the context the model just saw. + Before any call, falls back to a char-based estimate of ``messages``. + """ + if not self.last_request_usage.is_zero(): + return self.last_request_usage.prompt_tokens + return estimate_messages_tokens(self.messages) + + @property + def context_tokens_source(self) -> str: + """``api`` / ``estimate`` for :attr:`context_tokens`.""" + if not self.last_request_usage.is_zero(): + return self.last_request_usage.source + return "estimate" + def _ensure_system_prompt(self) -> None: """Prepend system prompt once when configured and history has none.""" if not self.system_prompt: diff --git a/src/plyngent/cli/repl.py b/src/plyngent/cli/repl.py index dc98ebf..dbfef49 100644 --- a/src/plyngent/cli/repl.py +++ b/src/plyngent/cli/repl.py @@ -54,28 +54,31 @@ _CONTENT_PREVIEW = 200 def _cmd_status(state: ReplState) -> None: - from plyngent.agent.budget import estimate_messages_chars, estimate_messages_tokens + from plyngent.agent.budget import estimate_messages_chars pending = state.agent.pending_retry_text pending_disp = "yes" if pending else "no" ctx_chars = estimate_messages_chars(state.agent.messages) - ctx_tokens = estimate_messages_tokens(state.agent.messages) + ctx_tokens = state.agent.context_tokens + ctx_src = state.agent.context_tokens_source 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 + # API prompt_tokens from the last model call is real context size for that request. + ctx_tag = "api" if ctx_src == "api" else "est" + ctx_tilde = "" if ctx_src == "api" else "~" 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, once) " + f"context_tokens={ctx_tilde}{ctx_tokens}/{ctx_budget} ({ctx_tag}) " f"context_chars={ctx_chars} " f"tool_result_max={state.agent.max_tool_result_chars}\n" - f"last_request={last_req.format_line()} " - f"(last model call; ~context size if from API)\n" + f"last_request={last_req.format_line()}\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" diff --git a/src/plyngent/cli/retry.py b/src/plyngent/cli/retry.py index 256de54..e31463d 100644 --- a/src/plyngent/cli/retry.py +++ b/src/plyngent/cli/retry.py @@ -94,8 +94,17 @@ def _echo_turn_usage(agent: ChatAgent) -> None: 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(): + # prompt_tokens on the last call ≈ context the model just saw + req = agent.last_request_usage + parts.append( + f"context={req.prompt_tokens} " + f"(prompt+completion={req.prompt_tokens}+{req.completion_tokens}" + f"={req.total_tokens}" + f"{' est' if req.source == 'estimate' else ''})" + ) + if not agent.last_turn_usage.is_zero() and ( + rounds > 1 or agent.last_turn_usage.total_tokens != agent.last_request_usage.total_tokens + ): label = agent.last_turn_usage.format_line(billed=True) if rounds > 1: parts.append(f"turn {label} over {rounds} rounds") diff --git a/tests/test_agent/test_loop.py b/tests/test_agent/test_loop.py index 23709ee..341b126 100644 --- a/tests/test_agent/test_loop.py +++ b/tests/test_agent/test_loop.py @@ -273,6 +273,21 @@ async def test_chat_agent_turn_usage_sums_tool_rounds() -> None: assert agent.last_request_usage.prompt_tokens == 200 assert agent.last_turn_usage.prompt_tokens == 300 assert agent.last_turn_usage.total_tokens == 308 + # Context size is last request prompt, not billed sum + assert agent.context_tokens == 200 + assert agent.context_tokens_source == "api" + + +async def test_context_tokens_falls_back_to_message_estimate() -> None: + agent = ChatAgent( + ScriptedClient([]), + model="m", + stream=False, + messages=[UserChatMessage(content="12345678")], + ) + assert agent.last_request_usage.is_zero() + assert agent.context_tokens_source == "estimate" + assert agent.context_tokens >= 1 async def test_stream_yields_deltas_incrementally() -> None: diff --git a/tests/test_cli/test_repl_commands.py b/tests/test_cli/test_repl_commands.py index 481262d..65b125d 100644 --- a/tests/test_cli/test_repl_commands.py +++ b/tests/test_cli/test_repl_commands.py @@ -145,12 +145,25 @@ async def test_rounds(state: ReplState) -> None: async def test_status_shows_context_tokens( state: ReplState, capsys: pytest.CaptureFixture[str] ) -> None: + from plyngent.agent.usage import TokenUsage from plyngent.lmproto.openai_compatible.model import UserChatMessage state.agent.messages = [UserChatMessage(content="hello")] assert await handle_slash(state, "/status") is True out = capsys.readouterr().out - assert "context_tokens~=" in out + assert "context_tokens=" in out + assert "(est)" in out # no API usage yet assert "context_chars=" in out assert "tool_result_max=" in out assert str(state.workspace) in out + + state.agent.last_request_usage = TokenUsage( + prompt_tokens=1234, + completion_tokens=10, + total_tokens=1244, + source="api", + ) + assert await handle_slash(state, "/status") is True + out2 = capsys.readouterr().out + assert "context_tokens=1234/" in out2 + assert "(api)" in out2