Grayscale/nanobot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(providers): define typed usage contract

Browse Source
This commit is contained in:
chengyongru
2026-08-25 01:04:25 +08:00
committed by chengyongru
parent 89c94d8744
commit 9895c23cb5
84 changed files with 1643 additions and 726 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nanobot/agent/hook.py
+6 -6
View File
@@ -9,7 +9,7 @@ from typing import Any
from loguru import logger
from nanobot.providers.base import LLMResponse, ToolCallRequest
from nanobot.providers.base import LLMResponse, LLMUsage, ToolCallRequest
@dataclass(slots=True)
@@ -19,7 +19,7 @@ class AgentHookContext:
iteration: int
messages: list[dict[str, Any]]
response: LLMResponse | None = None
usage: dict[str, int] = field(default_factory=dict)
usage: LLMUsage | None = None
tool_calls: list[ToolCallRequest] = field(default_factory=list)
tool_results: list[Any] = field(default_factory=list)
tool_events: list[dict[str, str]] = field(default_factory=list)
@@ -39,7 +39,7 @@ class AgentRunHookContext:
messages: list[dict[str, Any]]
final_content: str | None = None
tools_used: list[str] = field(default_factory=list)
usage: dict[str, int] = field(default_factory=dict)
usage: LLMUsage | None = None
stop_reason: str | None = None
error: str | None = None
tool_events: list[dict[str, str]] = field(default_factory=list)
@@ -284,7 +284,7 @@ class SDKCaptureHook(AgentHook):
super().__init__()
self.tools_used: list[str] = []
self.messages: list[dict[str, Any]] = []
self.usage: dict[str, int] = {}
self.usage: LLMUsage | None = None
self.stop_reason: str | None = None
self.error: str | None = None
self.tool_events: list[dict[str, str]] = []
@@ -294,7 +294,7 @@ class SDKCaptureHook(AgentHook):
for call in context.tool_calls:
self.tools_used.append(call.name)
self.messages = list(context.messages)
self.usage = dict(context.usage)
self.usage = context.usage
self.stop_reason = context.stop_reason
self.error = context.error
self.tool_events = list(context.tool_events)
@@ -302,7 +302,7 @@ class SDKCaptureHook(AgentHook):
async def after_run(self, context: AgentRunHookContext) -> None:
self.tools_used = list(context.tools_used)
self.messages = list(context.messages)
self.usage = dict(context.usage)
self.usage = context.usage
self.stop_reason = context.stop_reason
self.error = context.error
self.tool_events = list(context.tool_events)
nanobot/agent/loop.py
+7 -7
View File
@@ -49,7 +49,7 @@ from nanobot.bus.queue import MessageBus
from nanobot.bus.runtime_events import RuntimeEventBus
from nanobot.command import CommandContext, CommandRouter, register_builtin_commands
from nanobot.config.schema import AgentDefaults, ModelPresetConfig
from nanobot.providers.base import LLMProvider, ProviderConversationState
from nanobot.providers.base import LLMProvider, LLMUsage, ProviderConversationState
from nanobot.providers.factory import ProviderSnapshot
from nanobot.runtime_context import (
RUNTIME_CONTEXT_HISTORY_META,
@@ -167,7 +167,7 @@ class TurnContext:
turn_wall_started_at: float = field(default_factory=time.time)
visible_run_started_at: float | None = None
turn_latency_ms: int | None = None
usage: dict[str, int] = field(default_factory=dict)
usage: LLMUsage | None = None
def require_runtime(self) -> LLMRuntime:
"""Return the runtime established by the BUILD stage."""
@@ -203,7 +203,7 @@ class AgentLoop:
return self.tools.tool_names
@property
def last_usage(self) -> Mapping[str, int]:
def last_usage(self) -> LLMUsage | None:
"""Latest aggregate usage exposed through the runtime-control snapshot."""
return self._last_usage
@@ -378,7 +378,7 @@ class AgentLoop:
default_restrict_to_workspace=restrict_to_workspace,
)
self._start_time = time.time()
self._last_usage: dict[str, int] = {}
self._last_usage: LLMUsage | None = None
self._extra_hooks: list[AgentHook] = hooks or []
self._hook_factories: list[AgentTurnHookFactory] = hook_factories or []
@@ -2031,7 +2031,7 @@ class AgentLoop:
ctx.all_messages = all_msgs
ctx.stop_reason = stop_reason
ctx.had_injections = had_injections
ctx.usage = dict(self._last_usage)
ctx.usage = self._last_usage
ctx.delivery.record_usage(ctx.usage)
if ctx.kind is TurnKind.USER:
await turn_continuation.maybe_continue_turn(ctx)
@@ -2058,8 +2058,8 @@ class AgentLoop:
else ctx.turn_wall_started_at
)
ctx.turn_latency_ms = max(0, int((time.time() - latency_started_at) * 1000))
if ctx.usage and not ctx.ephemeral:
session.metadata["_last_usage"] = dict(ctx.usage)
if ctx.usage is not None and not ctx.ephemeral:
session.metadata["_last_usage"] = ctx.usage.to_dict()
self._save_turn(
session, ctx.all_messages, ctx.save_skip,
turn_latency_ms=ctx.turn_latency_ms,
nanobot/agent/progress_hook.py
+7 -5
View File
@@ -210,12 +210,14 @@ class AgentProgressHook(AgentHook):
tool_hint=False,
tool_events=tool_events,
)
u = context.usage or {}
u = context.usage
logger.debug(
"LLM usage: prompt={} completion={} cached={}",
u.get("prompt_tokens", 0),
u.get("completion_tokens", 0),
u.get("cached_tokens", 0),
"LLM usage: input={} output={} cache_read={} cache_write={} source={}",
u.input_tokens if u else 0,
u.output_tokens if u else 0,
u.cache_read_tokens if u else None,
u.cache_write_tokens if u else None,
u.source if u else "missing",
)
def finalize_content(self, context: AgentHookContext, content: str | None) -> str | None:
nanobot/agent/runner.py
+41 -78
View File
@@ -23,6 +23,7 @@ from nanobot.agent.tools.registry import ToolRegistry, is_tool_error_result
from nanobot.providers.base import (
LLMProvider,
LLMResponse,
LLMUsage,
ProviderCallContext,
ProviderConversationState,
ToolCallRequest,
@@ -126,7 +127,7 @@ class AgentRunResult:
final_content: str | None
messages: list[dict[str, Any]]
tools_used: list[str] = field(default_factory=list)
usage: dict[str, int] = field(default_factory=dict)
usage: LLMUsage | None = None
stop_reason: str = "completed"
error: str | None = None
tool_events: list[dict[str, str]] = field(default_factory=list)
@@ -412,7 +413,7 @@ class AgentRunner:
context.messages = deepcopy(result.messages)
context.final_content = result.final_content
context.tools_used = list(result.tools_used)
context.usage = dict(result.usage)
context.usage = result.usage
context.stop_reason = result.stop_reason
context.error = result.error
context.tool_events = deepcopy(result.tool_events)
@@ -443,7 +444,7 @@ class AgentRunner:
) -> AgentRunResult:
final_content: str | None = None
tools_used: list[str] = []
usage = {"prompt_tokens": 0, "completion_tokens": 0}
usage: LLMUsage | None = None
error: str | None = None
stop_reason = "completed"
tool_events: list[dict[str, str]] = []
@@ -519,8 +520,8 @@ class AgentRunner:
)
response.content = cleaned_content
raw_usage = self._usage_or_estimate(spec, messages_for_model, response)
context.usage = dict(raw_usage)
self._accumulate_usage(usage, raw_usage)
context.usage = raw_usage
usage = self._merge_usage(usage, raw_usage)
if reasoning_text and not context.streamed_reasoning:
await hook.emit_reasoning(reasoning_text)
await hook.emit_reasoning_end()
@@ -683,10 +684,10 @@ class AgentRunner:
conversation_state=conversation_state,
)
retry_usage = self._usage_or_estimate(spec, retry_messages, response)
self._accumulate_usage(usage, retry_usage)
usage = self._merge_usage(usage, retry_usage)
raw_usage = self._merge_usage(raw_usage, retry_usage)
context.response = response
context.usage = dict(raw_usage)
context.usage = raw_usage
context.tool_calls = list(response.tool_calls)
original_content = response.content
clean = hook.finalize_content(context, response.content)
@@ -859,7 +860,7 @@ class AgentRunner:
had_injections = True
terminal_content = None
if spec.finalize_on_max_iterations:
terminal_content = await self._try_finalize_after_max_iterations(
terminal_content, usage = await self._try_finalize_after_max_iterations(
spec,
hook,
messages,
@@ -1236,9 +1237,9 @@ class AgentRunner:
spec: AgentRunSpec,
hook: AgentHook,
messages: list[dict[str, Any]],
usage: dict[str, int],
usage: LLMUsage | None,
conversation_state: ProviderConversationStateController,
) -> str | None:
) -> tuple[str | None, LLMUsage | None]:
retry_messages = self._budget_exhausted_finalization_messages(messages)
try:
response = await self._request_no_tools(
@@ -1253,10 +1254,10 @@ class AgentRunner:
"Budget-exhausted finalization failed for {}; using fallback",
spec.session_key or "default",
)
return None
return None, usage
raw_usage = self._usage_or_estimate(spec, retry_messages, response)
self._accumulate_usage(usage, raw_usage)
usage = self._merge_usage(usage, raw_usage)
if response.finish_reason == "error" or response.has_tool_calls:
logger.warning(
"Budget-exhausted finalization returned finish_reason='{}' "
@@ -1265,19 +1266,19 @@ class AgentRunner:
len(response.tool_calls),
spec.session_key or "default",
)
return None
return None, usage
context = AgentHookContext(
iteration=spec.max_iterations,
messages=messages,
response=response,
usage=dict(raw_usage),
usage=raw_usage,
session_key=spec.session_key,
)
clean = hook.finalize_content(context, response.content)
if is_blank_text(clean):
return None
return clean
return None, usage
return clean, usage
async def _request_no_tools(
self,
@@ -1349,31 +1350,24 @@ class AgentRunner:
spec: AgentRunSpec,
messages: list[dict[str, Any]],
response: LLMResponse,
) -> dict[str, int]:
usage = self._usage_dict(response.usage)
total = self._usage_total(usage)
if total > 0:
usage["total_tokens"] = total
usage.setdefault("provider_tokens", total)
elif response.finish_reason == "error":
return {}
else:
) -> LLMUsage | None:
usage = response.usage
if response.finish_reason == "error":
if usage is None or usage.total_tokens == 0:
usage = LLMUsage.empty_request()
elif usage is None or usage.total_tokens == 0:
usage = self._estimate_response_usage(spec, messages, response)
completion = usage.get("completion_tokens", 0)
if response.generation_ms is not None and completion > 0:
usage["generation_ms"] = response.generation_ms
usage["measured_completion_tokens"] = completion
if response.ttft_ms is not None:
usage["ttft_ms"] = response.ttft_ms
usage["timed_requests"] = 1
return usage
return usage.with_timing(
generation_ms=response.generation_ms,
ttft_ms=response.ttft_ms,
)
def _estimate_response_usage(
self,
spec: AgentRunSpec,
messages: list[dict[str, Any]],
response: LLMResponse,
) -> dict[str, int]:
) -> LLMUsage:
try:
tools = spec.tools.get_definitions()
except Exception:
@@ -1391,52 +1385,21 @@ class AgentRunner:
thinking_blocks=response.thinking_blocks,
)
completion_tokens = estimate_message_tokens(assistant_message)
total_tokens = max(0, prompt_tokens) + max(0, completion_tokens)
if total_tokens <= 0:
return {}
return {
"prompt_tokens": max(0, prompt_tokens),
"completion_tokens": max(0, completion_tokens),
"total_tokens": total_tokens,
"estimated_tokens": total_tokens,
}
return LLMUsage.estimated(
input_tokens=max(0, prompt_tokens),
output_tokens=max(0, completion_tokens),
)
@staticmethod
def _usage_dict(usage: dict[str, Any] | None) -> dict[str, int]:
if not usage:
return {}
result: dict[str, int] = {}
for key, value in usage.items():
try:
result[key] = int(value or 0)
except (TypeError, ValueError):
continue
return result
@staticmethod
def _usage_total(usage: dict[str, int]) -> int:
return max(0, usage.get("total_tokens", 0) or (
usage.get("prompt_tokens", 0) + usage.get("completion_tokens", 0)
))
@staticmethod
def _merge_usage(left: dict[str, int], right: dict[str, int]) -> dict[str, int]:
merged = dict(left)
for key, value in right.items():
merged[key] = merged.get(key, 0) + value
return merged
@staticmethod
def _accumulate_usage(total: dict[str, int], request: dict[str, int]) -> None:
"""Fold one model request into the current turn's usage."""
total["request_count"] = total.get("request_count", 0) + 1
prompt_tokens = request.get("prompt_tokens")
if prompt_tokens is not None and prompt_tokens >= 0:
total["context_tokens"] = prompt_tokens
for key, value in request.items():
if key in {"context_tokens", "request_count"} or value < 0:
continue
total[key] = total.get(key, 0) + value
def _merge_usage(
left: LLMUsage | None,
right: LLMUsage | None,
) -> LLMUsage | None:
if left is None:
return right
if right is None:
return left
return left + right
async def _execute_tools(
self,
nanobot/agent/subagent.py
+3 -3
View File
@@ -28,7 +28,7 @@ from nanobot.agent.tools.registry import ToolRegistry
from nanobot.bus.events import InboundMessage
from nanobot.bus.queue import MessageBus
from nanobot.config.schema import AgentDefaults, ToolsConfig
from nanobot.providers.base import LLMProvider
from nanobot.providers.base import LLMProvider, LLMUsage
from nanobot.security.workspace_access import (
WorkspaceScope,
bind_workspace_scope,
@@ -56,7 +56,7 @@ class SubagentStatus:
phase: str = "initializing" # initializing | awaiting_tools | tools_completed | final_response | done | error
iteration: int = 0
tool_events: list[dict[str, str]] = field(default_factory=list)
usage: dict[str, int] = field(default_factory=dict)
usage: LLMUsage | None = None
stop_reason: str | None = None
error: str | None = None
@@ -82,7 +82,7 @@ class _SubagentHook(AgentHook):
return
self._status.iteration = context.iteration
self._status.tool_events = list(context.tool_events)
self._status.usage = dict(context.usage)
self._status.usage = context.usage
if context.error:
self._status.error = str(context.error)
nanobot/agent/tools/runtime_control.py
+5 -4
View File
@@ -12,6 +12,7 @@ if TYPE_CHECKING:
from nanobot.agent.tools.shell import ExecToolConfig
from nanobot.agent.tools.web import WebToolsConfig
from nanobot.config.schema import ModelPresetConfig
from nanobot.providers.base import LLMUsage
from nanobot.utils.llm_runtime import LLMRuntime
@@ -65,7 +66,7 @@ class RuntimeSnapshot:
web_config: dict[str, object]
exec_config: dict[str, object]
subagent_statuses: dict[str, dict[str, object]]
last_usage: dict[str, int]
last_usage: Mapping[str, JsonScalar]
scratchpad: dict[str, JsonValue]
def as_mapping(self) -> Mapping[str, object]:
@@ -151,7 +152,7 @@ class _RuntimeControlTarget(Protocol):
def tool_names(self) -> list[str]: ...
@property
def last_usage(self) -> Mapping[str, int]: ...
def last_usage(self) -> LLMUsage | None: ...
def set_runtime_model(self, model: str) -> LLMRuntime: ...
@@ -190,7 +191,7 @@ class AgentRuntimeControl:
web_config=_snapshot_web_config(target.web_config),
exec_config=_snapshot_exec_config(target.exec_config),
subagent_statuses=_snapshot_subagent_statuses(target.subagents),
last_usage=dict(target.last_usage),
last_usage=target.last_usage.to_dict() if target.last_usage is not None else {},
scratchpad=_snapshot_json_mapping(self.__scratchpad),
)
@@ -297,7 +298,7 @@ def _snapshot_subagent_status(status: SubagentStatus) -> dict[str, object]:
"phase": status.phase,
"iteration": status.iteration,
"tool_events": [dict(event) for event in status.tool_events],
"usage": dict(status.usage),
"usage": status.usage.to_dict() if status.usage is not None else None,
"stop_reason": status.stop_reason,
"error": status.error,
}
nanobot/agent/tools/self.py
+1 -1
View File
@@ -150,7 +150,7 @@ class MyTool(Tool):
"Actions: check, set.\n"
"- check (no key): full config overview — start here.\n"
"- check (key): drill into a value. Dot-paths allowed "
"(e.g. '_last_usage.prompt_tokens', 'web_config.enable').\n"
"(e.g. '_last_usage.input_tokens', 'web_config.enable').\n"
"- set (key, value): change config or store notes in your scratchpad. "
"Scratchpad keys persist across turns but not restarts.\n"
"Key values: _current_iteration (current progress), "
nanobot/agent/turn_delivery.py
+3 -2
View File
@@ -4,7 +4,7 @@ from __future__ import annotations
import dataclasses
import time
from collections.abc import Awaitable, Callable, Mapping
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, cast
@@ -19,6 +19,7 @@ from nanobot.bus.outbound_events import (
from nanobot.bus.progress import build_bus_progress_callback
from nanobot.bus.queue import MessageBus
from nanobot.bus.runtime_events import RuntimeEventBus, RuntimeEventPublisher
from nanobot.providers.base import LLMUsage
if TYPE_CHECKING:
from nanobot.utils.llm_runtime import LLMRuntime
@@ -203,7 +204,7 @@ class TurnDelivery:
def record_latency(self, latency_ms: int | None) -> None:
self.runtime_event_publisher.record_turn_latency(self.session_key, latency_ms)
def record_usage(self, usage: Mapping[str, int]) -> None:
def record_usage(self, usage: LLMUsage | None) -> None:
self.runtime_event_publisher.record_turn_usage(self.session_key, usage)
def background_response(