mirror of
https://github.com/NCBM/plyngent.git
synced 2026-07-23 05:55:16 +08:00
core/agent+cli: separate last-request usage from billed turn totals
Tool-loop rounds re-send history, so summed prompt tokens are billing not context size. Track last_request_usage and label turn/session as billed; clarify /status and end-of-turn lines.
This commit is contained in:
@@ -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/`)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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}"
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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}"
|
||||
)
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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."""
|
||||
|
||||
|
||||
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user