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()`.
|
- **`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.
|
- **`/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`).
|
- 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).
|
- Config ``[agent]``: `system_prompt`, `max_tool_result_chars`, `parallel_tools`, `confirm_destructive`, `path_denylist`, `max_context_tokens` (default 200k est. tokens).
|
||||||
|
|
||||||
### Tools (`tools/`)
|
### Tools (`tools/`)
|
||||||
|
|||||||
@@ -42,6 +42,8 @@ class ChatAgent:
|
|||||||
pending_retry_text: str | None
|
pending_retry_text: str | None
|
||||||
session_usage: TokenUsage
|
session_usage: TokenUsage
|
||||||
last_turn_usage: TokenUsage
|
last_turn_usage: TokenUsage
|
||||||
|
last_request_usage: TokenUsage
|
||||||
|
last_turn_rounds: int
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
@@ -78,6 +80,8 @@ class ChatAgent:
|
|||||||
self.pending_retry_text = None
|
self.pending_retry_text = None
|
||||||
self.session_usage = TokenUsage()
|
self.session_usage = TokenUsage()
|
||||||
self.last_turn_usage = TokenUsage()
|
self.last_turn_usage = TokenUsage()
|
||||||
|
self.last_request_usage = TokenUsage()
|
||||||
|
self.last_turn_rounds = 0
|
||||||
self._ensure_system_prompt()
|
self._ensure_system_prompt()
|
||||||
|
|
||||||
def _ensure_system_prompt(self) -> None:
|
def _ensure_system_prompt(self) -> None:
|
||||||
@@ -123,6 +127,8 @@ class ChatAgent:
|
|||||||
|
|
||||||
completed = False
|
completed = False
|
||||||
turn_usage = TokenUsage()
|
turn_usage = TokenUsage()
|
||||||
|
turn_rounds = 0
|
||||||
|
last_request = TokenUsage()
|
||||||
try:
|
try:
|
||||||
async for event in run_chat_loop(
|
async for event in run_chat_loop(
|
||||||
self.client,
|
self.client,
|
||||||
@@ -138,6 +144,9 @@ class ChatAgent:
|
|||||||
max_context_tokens=self.max_context_tokens,
|
max_context_tokens=self.max_context_tokens,
|
||||||
):
|
):
|
||||||
if isinstance(event, UsageEvent):
|
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)
|
turn_usage = turn_usage.add(event.usage)
|
||||||
self.session_usage = self.session_usage.add(event.usage)
|
self.session_usage = self.session_usage.add(event.usage)
|
||||||
yield event
|
yield event
|
||||||
@@ -148,6 +157,8 @@ class ChatAgent:
|
|||||||
raise
|
raise
|
||||||
|
|
||||||
self.last_turn_usage = turn_usage
|
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:]:
|
for message in self.messages[pre_len:]:
|
||||||
await self._persist(message)
|
await self._persist(message)
|
||||||
self.pending_retry_text = None
|
self.pending_retry_text = None
|
||||||
|
|||||||
@@ -40,15 +40,21 @@ class TokenUsage(Struct, omit_defaults=True):
|
|||||||
def is_zero(self) -> bool:
|
def is_zero(self) -> bool:
|
||||||
return self.prompt_tokens == 0 and self.completion_tokens == 0 and self.total_tokens == 0
|
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 = ""
|
tag = ""
|
||||||
if self.source == "estimate":
|
if self.source == "estimate":
|
||||||
tag = " (est)"
|
tag = " (est)"
|
||||||
elif self.source == "mixed":
|
elif self.source == "mixed":
|
||||||
tag = " (api+est)"
|
tag = " (api+est)"
|
||||||
|
prefix = "billed " if billed else ""
|
||||||
return (
|
return (
|
||||||
f"tokens prompt={self.prompt_tokens} completion={self.completion_tokens} "
|
f"{prefix}tokens prompt={self.prompt_tokens} "
|
||||||
f"total={self.total_tokens}{tag}"
|
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
|
ctx_budget = state.agent.max_context_tokens
|
||||||
session_u = state.agent.session_usage
|
session_u = state.agent.session_usage
|
||||||
last_u = state.agent.last_turn_usage
|
last_u = state.agent.last_turn_usage
|
||||||
|
last_req = state.agent.last_request_usage
|
||||||
|
last_rounds = state.agent.last_turn_rounds
|
||||||
click.echo(
|
click.echo(
|
||||||
f"provider={state.provider_name} model={state.model}\n"
|
f"provider={state.provider_name} model={state.model}\n"
|
||||||
f"session={state.session_id} messages={len(state.agent.messages)} "
|
f"session={state.session_id} messages={len(state.agent.messages)} "
|
||||||
f"pending_retry={pending_disp}\n"
|
f"pending_retry={pending_disp}\n"
|
||||||
f"tools={'on' if state.tools_enabled else 'off'} "
|
f"tools={'on' if state.tools_enabled else 'off'} "
|
||||||
f"rounds={state.max_rounds} stream={'on' if state.agent.stream else 'off'}\n"
|
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"context_chars={ctx_chars} "
|
||||||
f"tool_result_max={state.agent.max_tool_result_chars}\n"
|
f"tool_result_max={state.agent.max_tool_result_chars}\n"
|
||||||
f"usage_session={session_u.format_line()}\n"
|
f"last_request={last_req.format_line()} "
|
||||||
f"usage_last_turn={last_u.format_line()}\n"
|
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}"
|
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:
|
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():
|
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:
|
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)
|
agent = ChatAgent(client, model="m", stream=False)
|
||||||
_ = [e async for e in agent.run("one")]
|
_ = [e async for e in agent.run("one")]
|
||||||
assert agent.last_turn_usage.total_tokens == 6
|
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
|
assert agent.session_usage.total_tokens == 6
|
||||||
_ = [e async for e in agent.run("two")]
|
_ = [e async for e in agent.run("two")]
|
||||||
assert agent.last_turn_usage.total_tokens == 10
|
assert agent.last_turn_usage.total_tokens == 10
|
||||||
|
assert agent.last_request_usage.total_tokens == 10
|
||||||
assert agent.session_usage.total_tokens == 16
|
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:
|
async def test_stream_yields_deltas_incrementally() -> None:
|
||||||
"""Text deltas are yielded as chunks arrive, not only after the full stream."""
|
"""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
|
assert "(est)" in line
|
||||||
mixed = TokenUsage(prompt_tokens=1, completion_tokens=0, total_tokens=1, source="mixed").format_line()
|
mixed = TokenUsage(prompt_tokens=1, completion_tokens=0, total_tokens=1, source="mixed").format_line()
|
||||||
assert "(api+est)" in mixed
|
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