mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-31 08:13:11 +03:00
feat(usage): add unified provider usage backend
This commit is contained in:
@@ -49,6 +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.llm_usage.context import source_from_request
|
||||
from nanobot.providers.base import LLMProvider, LLMUsage, ProviderConversationState
|
||||
from nanobot.providers.factory import ProviderSnapshot
|
||||
from nanobot.runtime_context import (
|
||||
@@ -1202,6 +1203,11 @@ class AgentLoop:
|
||||
message_metadata=metadata,
|
||||
),
|
||||
provider_state=provider_state,
|
||||
llm_usage_source=source_from_request(
|
||||
active_session_key,
|
||||
channel=channel,
|
||||
metadata=metadata,
|
||||
),
|
||||
))
|
||||
finally:
|
||||
turn_scope_stack.close()
|
||||
|
||||
+11
-9
@@ -20,6 +20,7 @@ from typing import TYPE_CHECKING, Any, Callable, Iterator, cast
|
||||
|
||||
from loguru import logger
|
||||
|
||||
from nanobot.llm_usage.context import llm_usage_source
|
||||
from nanobot.runtime_context import public_history_messages
|
||||
from nanobot.session.manager import (
|
||||
MIN_COMPACTED_REPLAY_MESSAGES,
|
||||
@@ -915,15 +916,16 @@ class Consolidator:
|
||||
if not messages:
|
||||
return None
|
||||
try:
|
||||
response = await runtime.provider.chat_with_retry(
|
||||
model=runtime.model,
|
||||
messages=request_messages,
|
||||
tools=request_tools,
|
||||
tool_choice="none",
|
||||
temperature=runtime.generation.temperature,
|
||||
max_tokens=runtime.generation.max_tokens,
|
||||
reasoning_effort=runtime.generation.reasoning_effort,
|
||||
)
|
||||
with llm_usage_source("dream"):
|
||||
response = await runtime.provider.chat_with_retry(
|
||||
model=runtime.model,
|
||||
messages=request_messages,
|
||||
tools=request_tools,
|
||||
tool_choice="none",
|
||||
temperature=runtime.generation.temperature,
|
||||
max_tokens=runtime.generation.max_tokens,
|
||||
reasoning_effort=runtime.generation.reasoning_effort,
|
||||
)
|
||||
except Exception:
|
||||
logger.warning("Consolidation provider call failed, raw-dumping to history")
|
||||
self.store.raw_archive(messages, session_key=session_key)
|
||||
|
||||
+23
-10
@@ -20,6 +20,12 @@ from nanobot.agent.context_governance import (
|
||||
)
|
||||
from nanobot.agent.hook import AgentHook, AgentHookContext, AgentRunHookContext
|
||||
from nanobot.agent.tools.registry import ToolRegistry, is_tool_error_result
|
||||
from nanobot.llm_usage.context import (
|
||||
LLMUsageSource,
|
||||
bind_llm_usage_source,
|
||||
reset_llm_usage_source,
|
||||
source_from_session_key,
|
||||
)
|
||||
from nanobot.providers.base import (
|
||||
LLMProvider,
|
||||
LLMResponse,
|
||||
@@ -118,6 +124,7 @@ class AgentRunSpec:
|
||||
goal_continue_message: GoalContinueMessage | None = None
|
||||
finalize_on_max_iterations: bool = True
|
||||
provider_state: ProviderConversationState | None = None
|
||||
llm_usage_source: LLMUsageSource | None = None
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
@@ -392,6 +399,9 @@ class AgentRunner:
|
||||
hook = spec.hook or AgentHook()
|
||||
messages = list(spec.initial_messages)
|
||||
context = AgentRunHookContext(messages=deepcopy(messages))
|
||||
llm_usage_source_token = bind_llm_usage_source(
|
||||
spec.llm_usage_source or source_from_session_key(spec.session_key)
|
||||
)
|
||||
|
||||
try:
|
||||
await hook.before_run(context)
|
||||
@@ -424,17 +434,20 @@ class AgentRunner:
|
||||
await hook.after_run(context)
|
||||
return result
|
||||
finally:
|
||||
context.messages = deepcopy(messages)
|
||||
if context.exception is None:
|
||||
await hook.on_finally(context)
|
||||
else:
|
||||
try:
|
||||
try:
|
||||
context.messages = deepcopy(messages)
|
||||
if context.exception is None:
|
||||
await hook.on_finally(context)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"AgentHook.on_finally error after {}",
|
||||
context.stop_reason or "run exception",
|
||||
)
|
||||
else:
|
||||
try:
|
||||
await hook.on_finally(context)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"AgentHook.on_finally error after {}",
|
||||
context.stop_reason or "run exception",
|
||||
)
|
||||
finally:
|
||||
reset_llm_usage_source(llm_usage_source_token)
|
||||
|
||||
async def _run_core(
|
||||
self,
|
||||
|
||||
@@ -8,7 +8,7 @@ import warnings
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, TypedDict
|
||||
from typing import Any, Callable, NotRequired, TypedDict
|
||||
|
||||
from loguru import logger
|
||||
|
||||
@@ -28,6 +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.llm_usage.context import LLMUsageSource, current_llm_usage_source
|
||||
from nanobot.providers.base import LLMProvider, LLMUsage
|
||||
from nanobot.security.workspace_access import (
|
||||
WorkspaceScope,
|
||||
@@ -43,6 +44,7 @@ class _SubagentOrigin(TypedDict):
|
||||
channel: str
|
||||
chat_id: str
|
||||
session_key: str | None
|
||||
llm_usage_source: NotRequired[LLMUsageSource]
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
@@ -252,6 +254,7 @@ class SubagentManager:
|
||||
"channel": origin_channel,
|
||||
"chat_id": origin_chat_id,
|
||||
"session_key": session_key,
|
||||
"llm_usage_source": current_llm_usage_source(),
|
||||
}
|
||||
|
||||
status = SubagentStatus(
|
||||
@@ -315,6 +318,7 @@ class SubagentManager:
|
||||
"channel": origin_channel,
|
||||
"chat_id": origin_chat_id,
|
||||
"session_key": session_key,
|
||||
"llm_usage_source": current_llm_usage_source(),
|
||||
}
|
||||
status = SubagentStatus(
|
||||
task_id=task_id,
|
||||
@@ -417,6 +421,10 @@ class SubagentManager:
|
||||
session_key=sess_key,
|
||||
workspace=root,
|
||||
llm_timeout_s=llm_timeout,
|
||||
llm_usage_source=origin.get(
|
||||
"llm_usage_source",
|
||||
current_llm_usage_source(),
|
||||
),
|
||||
))
|
||||
finally:
|
||||
if token is not None:
|
||||
|
||||
@@ -313,6 +313,8 @@ def _run_gateway(
|
||||
from nanobot.cron.service import CronJobSkippedError, CronService
|
||||
from nanobot.cron.session_turns import is_bound_cron_job
|
||||
from nanobot.cron.types import CronJob
|
||||
from nanobot.llm_usage import record_llm_call
|
||||
from nanobot.llm_usage.context import llm_usage_source
|
||||
from nanobot.providers.factory import (
|
||||
ProviderSnapshot,
|
||||
build_provider_snapshot,
|
||||
@@ -330,7 +332,6 @@ def _run_gateway(
|
||||
)
|
||||
from nanobot.triggers.local_runner import run_local_trigger_queue
|
||||
from nanobot.triggers.local_store import LocalTriggerStore
|
||||
from nanobot.webui.token_usage import TokenUsageHook
|
||||
|
||||
port = port if port is not None else config.gateway.port
|
||||
webui_url = _webui_browser_url(config)
|
||||
@@ -361,7 +362,8 @@ def _run_gateway(
|
||||
runtime_events = RuntimeEventBus()
|
||||
fallback_model_observer = build_webui_fallback_model_observer(bus)
|
||||
|
||||
def _observe_fallback_models(snapshot: ProviderSnapshot) -> ProviderSnapshot:
|
||||
def _observe_provider(snapshot: ProviderSnapshot) -> ProviderSnapshot:
|
||||
snapshot.provider.set_llm_call_observer(record_llm_call)
|
||||
if isinstance(snapshot.provider, FallbackProvider):
|
||||
snapshot.provider.set_fallback_model_observer(fallback_model_observer)
|
||||
return snapshot
|
||||
@@ -371,20 +373,19 @@ def _run_gateway(
|
||||
**kwargs: Any,
|
||||
) -> ProviderSnapshot:
|
||||
try:
|
||||
return _observe_fallback_models(load_provider_snapshot(*args, **kwargs))
|
||||
return _observe_provider(load_provider_snapshot(*args, **kwargs))
|
||||
except ValueError as exc:
|
||||
if unconfigured_provider_error is None:
|
||||
raise
|
||||
return build_unconfigured_provider_snapshot(config, str(exc))
|
||||
return _observe_provider(build_unconfigured_provider_snapshot(config, str(exc)))
|
||||
|
||||
if unconfigured_provider_error is not None:
|
||||
provider_snapshot = build_unconfigured_provider_snapshot(
|
||||
config,
|
||||
unconfigured_provider_error,
|
||||
provider_snapshot = _observe_provider(
|
||||
build_unconfigured_provider_snapshot(config, unconfigured_provider_error)
|
||||
)
|
||||
else:
|
||||
try:
|
||||
provider_snapshot = _observe_fallback_models(build_provider_snapshot(config))
|
||||
provider_snapshot = _observe_provider(build_provider_snapshot(config))
|
||||
except ValueError as exc:
|
||||
console.print(f"[red]Error: {exc}[/red]")
|
||||
raise typer.Exit(1) from exc
|
||||
@@ -443,7 +444,6 @@ def _run_gateway(
|
||||
runtime_events=runtime_events,
|
||||
turn_delivery_factory=turn_delivery_factory,
|
||||
provider_signature=provider_snapshot.signature,
|
||||
hooks=[TokenUsageHook(timezone_name=config.agents.defaults.timezone)],
|
||||
local_trigger_store=trigger_store,
|
||||
hook_factories=[create_file_edit_activity_hook],
|
||||
tool_registry=tools,
|
||||
@@ -564,13 +564,6 @@ def _run_gateway(
|
||||
except Exception:
|
||||
logger.exception("Dream cron job failed")
|
||||
finally:
|
||||
from nanobot.webui.token_usage import record_response_token_usage
|
||||
|
||||
record_response_token_usage(
|
||||
resp,
|
||||
source="dream",
|
||||
timezone_name=config.agents.defaults.timezone,
|
||||
)
|
||||
sha = _commit_dream_changes(store)
|
||||
if sha:
|
||||
logger.info("Dream commit: {}", sha)
|
||||
@@ -630,14 +623,15 @@ def _run_gateway(
|
||||
evaluator_prompt = resolve_evaluator_prompt(config.workspace_path)
|
||||
|
||||
# Fail closed: stay silent on evaluator failure instead of notifying.
|
||||
should_notify = await evaluate_response(
|
||||
response=response,
|
||||
task_context=prompt,
|
||||
provider=agent.provider,
|
||||
model=agent.model,
|
||||
evaluator_prompt=evaluator_prompt,
|
||||
default_notify=False,
|
||||
)
|
||||
with llm_usage_source("cron"):
|
||||
should_notify = await evaluate_response(
|
||||
response=response,
|
||||
task_context=prompt,
|
||||
provider=agent.provider,
|
||||
model=agent.model,
|
||||
evaluator_prompt=evaluator_prompt,
|
||||
default_notify=False,
|
||||
)
|
||||
|
||||
if should_notify:
|
||||
logger.info("Heartbeat: completed, delivering response")
|
||||
|
||||
@@ -479,13 +479,6 @@ async def cmd_dream(ctx: CommandContext) -> OutboundMessage:
|
||||
elapsed = time.monotonic() - t0
|
||||
content = f"Dream failed after {elapsed:.1f}s: {e}"
|
||||
finally:
|
||||
from nanobot.webui.token_usage import record_response_token_usage
|
||||
|
||||
record_response_token_usage(
|
||||
resp,
|
||||
source="dream",
|
||||
timezone_name=getattr(loop.context, "timezone", None),
|
||||
)
|
||||
if store.git.is_initialized():
|
||||
commit_msg = build_dream_commit_message("dream: manual run", diff_body)
|
||||
sha = store.git.auto_commit(commit_msg)
|
||||
|
||||
@@ -0,0 +1,86 @@
|
||||
"""Unified, content-free LLM usage backend."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from loguru import logger
|
||||
|
||||
from nanobot.config.paths import get_data_dir
|
||||
from nanobot.llm_usage.models import LLMCallRecord
|
||||
from nanobot.llm_usage.store import LLMUsageStore
|
||||
|
||||
_STORES_LOCK = threading.Lock()
|
||||
_STORES: dict[Path, LLMUsageStore] = {}
|
||||
|
||||
|
||||
def empty_usage_payload() -> dict[str, Any]:
|
||||
return {
|
||||
"days": [],
|
||||
"total_tokens": 0,
|
||||
"total_tokens_30d": 0,
|
||||
"total_tokens_365d": 0,
|
||||
"reported_tokens_30d": 0,
|
||||
"estimated_tokens_30d": 0,
|
||||
"cache_read_tokens_30d": 0,
|
||||
"cache_read_observed_input_tokens_30d": 0,
|
||||
"cache_read_rate_30d": None,
|
||||
"peak_day_tokens": 0,
|
||||
"current_streak_days": 0,
|
||||
"longest_streak_days": 0,
|
||||
"active_days_30d": 0,
|
||||
"requests_30d": 0,
|
||||
"failed_requests_30d": 0,
|
||||
"providers_30d": [],
|
||||
"updated_at": None,
|
||||
}
|
||||
|
||||
|
||||
def llm_usage_store_path() -> Path:
|
||||
return get_data_dir() / "llm_usage.sqlite3"
|
||||
|
||||
|
||||
def get_llm_usage_store(path: Path | None = None) -> LLMUsageStore:
|
||||
resolved = (path or llm_usage_store_path()).resolve(strict=False)
|
||||
with _STORES_LOCK:
|
||||
store = _STORES.get(resolved)
|
||||
if store is None:
|
||||
store = LLMUsageStore(resolved)
|
||||
_STORES[resolved] = store
|
||||
return store
|
||||
|
||||
|
||||
def record_llm_call(call: LLMCallRecord) -> None:
|
||||
"""Default fail-open callback attached to gateway provider snapshots."""
|
||||
try:
|
||||
get_llm_usage_store().record(call)
|
||||
except Exception:
|
||||
logger.exception("failed to record LLM usage")
|
||||
|
||||
|
||||
def llm_usage_payload(
|
||||
*,
|
||||
days: int = 371,
|
||||
timezone_name: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
try:
|
||||
return get_llm_usage_store().usage_payload(
|
||||
days=days,
|
||||
timezone_name=timezone_name,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("failed to query LLM usage")
|
||||
return empty_usage_payload()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"LLMCallRecord",
|
||||
"LLMUsageStore",
|
||||
"empty_usage_payload",
|
||||
"get_llm_usage_store",
|
||||
"record_llm_call",
|
||||
"llm_usage_store_path",
|
||||
"llm_usage_payload",
|
||||
]
|
||||
@@ -0,0 +1,70 @@
|
||||
"""Request-local metadata for LLM usage records."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Generator, Mapping
|
||||
from contextlib import contextmanager
|
||||
from contextvars import ContextVar, Token
|
||||
from typing import Literal
|
||||
|
||||
LLMUsageSource = Literal["user", "api", "cron", "dream", "system"]
|
||||
|
||||
_CURRENT_SOURCE: ContextVar[LLMUsageSource] = ContextVar(
|
||||
"nanobot_llm_usage_source",
|
||||
default="system",
|
||||
)
|
||||
|
||||
|
||||
def source_from_session_key(session_key: str | None) -> LLMUsageSource:
|
||||
"""Classify a private session key without persisting that key."""
|
||||
key = session_key or ""
|
||||
if key.startswith("dream:"):
|
||||
return "dream"
|
||||
if key == "heartbeat" or key.startswith("cron:"):
|
||||
return "cron"
|
||||
if key.startswith("api:"):
|
||||
return "api"
|
||||
if key.startswith("system:"):
|
||||
return "system"
|
||||
return "user"
|
||||
|
||||
|
||||
def source_from_request(
|
||||
session_key: str | None,
|
||||
*,
|
||||
channel: str | None,
|
||||
metadata: Mapping[str, object] | None,
|
||||
) -> LLMUsageSource:
|
||||
"""Classify a turn from trusted ingress metadata without retaining identifiers."""
|
||||
values = metadata or {}
|
||||
if isinstance(values.get("_cron_trigger"), Mapping):
|
||||
return "cron"
|
||||
if isinstance(values.get("_local_trigger"), Mapping):
|
||||
return "cron"
|
||||
if channel == "api":
|
||||
return "api"
|
||||
if channel == "system":
|
||||
return "system"
|
||||
return source_from_session_key(session_key)
|
||||
|
||||
|
||||
def current_llm_usage_source() -> LLMUsageSource:
|
||||
return _CURRENT_SOURCE.get()
|
||||
|
||||
|
||||
def bind_llm_usage_source(source: LLMUsageSource) -> Token[LLMUsageSource]:
|
||||
return _CURRENT_SOURCE.set(source)
|
||||
|
||||
|
||||
def reset_llm_usage_source(token: Token[LLMUsageSource]) -> None:
|
||||
_CURRENT_SOURCE.reset(token)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def llm_usage_source(source: LLMUsageSource) -> Generator[None]:
|
||||
"""Bind a coarse usage source for nested provider calls."""
|
||||
token = bind_llm_usage_source(source)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
reset_llm_usage_source(token)
|
||||
@@ -0,0 +1,38 @@
|
||||
"""Content-free records emitted for physical LLM provider calls."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from nanobot.llm_usage.context import LLMUsageSource
|
||||
from nanobot.providers.base import LLMUsage
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class LLMCallRecord:
|
||||
"""The small, chart-oriented result of one provider call attempt.
|
||||
|
||||
Request messages, response text, reasoning, and tool payloads deliberately do
|
||||
not belong to this contract. Sessions already own that content.
|
||||
"""
|
||||
|
||||
started_at_ms: int
|
||||
duration_ms: int
|
||||
provider: str
|
||||
model: str
|
||||
source: LLMUsageSource
|
||||
stream: bool
|
||||
finish_reason: str
|
||||
usage: LLMUsage | None = None
|
||||
error_status_code: int | None = None
|
||||
error_kind: str | None = None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.started_at_ms < 0 or self.duration_ms < 0:
|
||||
raise ValueError("LLM usage timestamps must be non-negative")
|
||||
if not self.provider.strip() or not self.model.strip():
|
||||
raise ValueError("LLM usage provider and model must be non-empty")
|
||||
if self.source not in {"user", "api", "cron", "dream", "system"}:
|
||||
raise ValueError("invalid LLM usage source")
|
||||
if not self.finish_reason.strip():
|
||||
raise ValueError("LLM usage finish_reason must be non-empty")
|
||||
@@ -0,0 +1,560 @@
|
||||
"""SQLite persistence and chart queries for LLM usage records."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sqlite3
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Iterable
|
||||
from copy import deepcopy
|
||||
from datetime import date, datetime, timedelta, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any, cast
|
||||
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
|
||||
|
||||
from nanobot.llm_usage.models import LLMCallRecord
|
||||
|
||||
SCHEMA_VERSION = 1
|
||||
MAX_DAYS_RETAINED = 400
|
||||
MAX_CALLS_RETAINED = 100_000
|
||||
|
||||
_ERROR_KINDS = frozenset({
|
||||
"authentication",
|
||||
"cancelled",
|
||||
"configuration",
|
||||
"connection",
|
||||
"content_filter",
|
||||
"context_length",
|
||||
"empty",
|
||||
"http",
|
||||
"invalid_request",
|
||||
"overloaded",
|
||||
"permission",
|
||||
"rate_limit",
|
||||
"refusal",
|
||||
"server_error",
|
||||
"timeout",
|
||||
})
|
||||
_FINISH_REASONS = frozenset({
|
||||
"cancelled",
|
||||
"content_filter",
|
||||
"error",
|
||||
"function_call",
|
||||
"length",
|
||||
"refusal",
|
||||
"stop",
|
||||
"tool_calls",
|
||||
})
|
||||
|
||||
_USAGE_COLUMNS = (
|
||||
"input_tokens",
|
||||
"output_tokens",
|
||||
"cache_read_tokens",
|
||||
"cache_write_tokens",
|
||||
"cache_read_observed_input_tokens",
|
||||
"cache_write_observed_input_tokens",
|
||||
"total_tokens",
|
||||
"reported_tokens",
|
||||
"estimated_tokens",
|
||||
"generation_ms",
|
||||
"measured_output_tokens",
|
||||
"ttft_ms",
|
||||
"timed_requests",
|
||||
)
|
||||
_REQUEST_COLUMNS = (
|
||||
"requests",
|
||||
"successful_requests",
|
||||
"failed_requests",
|
||||
"reported_requests",
|
||||
"estimated_requests",
|
||||
)
|
||||
_AGGREGATE_SQL = """
|
||||
COALESCE(SUM(input_tokens), 0) AS input_tokens,
|
||||
COALESCE(SUM(output_tokens), 0) AS output_tokens,
|
||||
COALESCE(SUM(cache_read_tokens), 0) AS cache_read_tokens,
|
||||
COALESCE(SUM(cache_write_tokens), 0) AS cache_write_tokens,
|
||||
COALESCE(SUM(
|
||||
CASE WHEN cache_read_tokens IS NOT NULL THEN input_tokens ELSE 0 END
|
||||
), 0) AS cache_read_observed_input_tokens,
|
||||
COALESCE(SUM(
|
||||
CASE WHEN cache_write_tokens IS NOT NULL THEN input_tokens ELSE 0 END
|
||||
), 0) AS cache_write_observed_input_tokens,
|
||||
COALESCE(SUM(total_tokens), 0) AS total_tokens,
|
||||
COALESCE(SUM(reported_tokens), 0) AS reported_tokens,
|
||||
COALESCE(SUM(estimated_tokens), 0) AS estimated_tokens,
|
||||
COALESCE(SUM(generation_ms), 0) AS generation_ms,
|
||||
COALESCE(SUM(measured_output_tokens), 0) AS measured_output_tokens,
|
||||
COALESCE(SUM(ttft_ms), 0) AS ttft_ms,
|
||||
COALESCE(SUM(timed_requests), 0) AS timed_requests,
|
||||
COUNT(*) AS requests,
|
||||
COALESCE(SUM(CASE WHEN finish_reason IN ('error', 'cancelled') THEN 0 ELSE 1 END), 0)
|
||||
AS successful_requests,
|
||||
COALESCE(SUM(CASE WHEN finish_reason IN ('error', 'cancelled') THEN 1 ELSE 0 END), 0)
|
||||
AS failed_requests,
|
||||
COALESCE(SUM(
|
||||
CASE WHEN total_tokens IS NOT NULL AND NOT (
|
||||
estimated_tokens > 0 AND reported_tokens = 0
|
||||
) THEN 1 ELSE 0 END
|
||||
), 0) AS reported_requests,
|
||||
COALESCE(SUM(
|
||||
CASE WHEN estimated_tokens > 0 AND reported_tokens = 0 THEN 1 ELSE 0 END
|
||||
), 0) AS estimated_requests,
|
||||
COALESCE(SUM(duration_ms), 0) AS duration_ms
|
||||
"""
|
||||
|
||||
|
||||
def _zone(timezone_name: str | None) -> timezone | ZoneInfo:
|
||||
if not timezone_name:
|
||||
return timezone.utc
|
||||
try:
|
||||
return ZoneInfo(timezone_name)
|
||||
except ZoneInfoNotFoundError:
|
||||
return timezone.utc
|
||||
|
||||
|
||||
def _clean_error_kind(value: str | None) -> str | None:
|
||||
if value is None:
|
||||
return None
|
||||
cleaned = value.strip().lower()
|
||||
if not cleaned:
|
||||
return None
|
||||
return cleaned if cleaned in _ERROR_KINDS else "other"
|
||||
|
||||
|
||||
def _clean_finish_reason(value: str) -> str:
|
||||
cleaned = value.strip().lower()
|
||||
return cleaned if cleaned in _FINISH_REASONS else "other"
|
||||
|
||||
|
||||
def _clean_status_code(value: int | None) -> int | None:
|
||||
if value is None:
|
||||
return None
|
||||
try:
|
||||
status = int(value)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
return status if 100 <= status <= 599 else None
|
||||
|
||||
|
||||
def _as_int_row(row: sqlite3.Row) -> dict[str, int]:
|
||||
return {
|
||||
key: max(0, int(row[key] or 0))
|
||||
for key in (*_USAGE_COLUMNS, *_REQUEST_COLUMNS, "duration_ms")
|
||||
}
|
||||
|
||||
|
||||
def _empty_totals() -> dict[str, int]:
|
||||
return {key: 0 for key in (*_USAGE_COLUMNS, *_REQUEST_COLUMNS, "duration_ms")}
|
||||
|
||||
|
||||
def _sum_rows(rows: Iterable[dict[str, Any]]) -> dict[str, int]:
|
||||
totals = _empty_totals()
|
||||
for row in rows:
|
||||
for key in totals:
|
||||
totals[key] += max(0, int(row.get(key) or 0))
|
||||
return totals
|
||||
|
||||
|
||||
class LLMUsageStore:
|
||||
"""A small synchronous WAL database shared by gateway threads/processes."""
|
||||
|
||||
def __init__(self, path: Path) -> None:
|
||||
self.path = path
|
||||
self._lock = threading.RLock()
|
||||
self._connection: sqlite3.Connection | None = None
|
||||
self._connection_pid: int | None = None
|
||||
self._last_prune_utc_day: int | None = None
|
||||
self._writes_since_size_prune = 0
|
||||
self._write_version = 0
|
||||
self._cached_payload_key: tuple[int, str, str, int, int] | None = None
|
||||
self._cached_payload: dict[str, Any] | None = None
|
||||
|
||||
def _connect(self) -> sqlite3.Connection:
|
||||
pid = os.getpid()
|
||||
if self._connection is not None and self._connection_pid == pid:
|
||||
return self._connection
|
||||
if self._connection is not None:
|
||||
self._connection.close()
|
||||
self._cached_payload_key = None
|
||||
self._cached_payload = None
|
||||
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||||
connection = sqlite3.connect(
|
||||
self.path,
|
||||
timeout=0.25,
|
||||
isolation_level=None,
|
||||
check_same_thread=False,
|
||||
)
|
||||
connection.row_factory = sqlite3.Row
|
||||
connection.execute("PRAGMA busy_timeout = 250")
|
||||
connection.execute("PRAGMA journal_mode = WAL")
|
||||
connection.execute("PRAGMA synchronous = NORMAL")
|
||||
connection.execute("PRAGMA temp_store = MEMORY")
|
||||
connection.create_function("llm_usage_local_day", 2, self._local_day, deterministic=True)
|
||||
connection.executescript(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS llm_calls (
|
||||
id INTEGER PRIMARY KEY,
|
||||
started_at_ms INTEGER NOT NULL,
|
||||
duration_ms INTEGER NOT NULL,
|
||||
provider TEXT NOT NULL,
|
||||
model TEXT NOT NULL,
|
||||
source TEXT NOT NULL,
|
||||
stream INTEGER NOT NULL,
|
||||
finish_reason TEXT NOT NULL,
|
||||
input_tokens INTEGER,
|
||||
output_tokens INTEGER,
|
||||
total_tokens INTEGER,
|
||||
cache_read_tokens INTEGER,
|
||||
cache_write_tokens INTEGER,
|
||||
reported_tokens INTEGER,
|
||||
estimated_tokens INTEGER,
|
||||
generation_ms INTEGER,
|
||||
measured_output_tokens INTEGER,
|
||||
ttft_ms INTEGER,
|
||||
timed_requests INTEGER,
|
||||
error_status_code INTEGER,
|
||||
error_kind TEXT
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS llm_calls_started_at_idx
|
||||
ON llm_calls(started_at_ms);
|
||||
CREATE INDEX IF NOT EXISTS llm_calls_provider_model_time_idx
|
||||
ON llm_calls(provider, model, started_at_ms);
|
||||
"""
|
||||
)
|
||||
connection.execute(f"PRAGMA user_version = {SCHEMA_VERSION}")
|
||||
self._connection = connection
|
||||
self._connection_pid = pid
|
||||
return connection
|
||||
|
||||
def _read_connection(self) -> sqlite3.Connection:
|
||||
connection = sqlite3.connect(
|
||||
self.path,
|
||||
timeout=0.25,
|
||||
isolation_level=None,
|
||||
)
|
||||
connection.row_factory = sqlite3.Row
|
||||
connection.execute("PRAGMA busy_timeout = 250")
|
||||
connection.execute("PRAGMA query_only = ON")
|
||||
connection.execute("PRAGMA temp_store = MEMORY")
|
||||
connection.create_function("llm_usage_local_day", 2, self._local_day, deterministic=True)
|
||||
return connection
|
||||
|
||||
@staticmethod
|
||||
def _local_day(started_at_ms: object, timezone_name: object) -> str | None:
|
||||
if not isinstance(started_at_ms, int) or not isinstance(timezone_name, str):
|
||||
return None
|
||||
dt = datetime.fromtimestamp(started_at_ms / 1000, timezone.utc)
|
||||
return dt.astimezone(_zone(timezone_name)).date().isoformat()
|
||||
|
||||
def close(self) -> None:
|
||||
with self._lock:
|
||||
if self._connection is not None:
|
||||
self._connection.close()
|
||||
self._connection = None
|
||||
self._connection_pid = None
|
||||
self._cached_payload_key = None
|
||||
self._cached_payload = None
|
||||
|
||||
def record(self, call: LLMCallRecord) -> None:
|
||||
usage = call.usage
|
||||
usage_data = usage.to_dict() if usage is not None else {}
|
||||
values: tuple[object, ...] = (
|
||||
call.started_at_ms,
|
||||
call.duration_ms,
|
||||
call.provider[:120],
|
||||
call.model[:240],
|
||||
call.source,
|
||||
int(call.stream),
|
||||
_clean_finish_reason(call.finish_reason),
|
||||
*(
|
||||
usage_data.get(key)
|
||||
for key in (
|
||||
"input_tokens",
|
||||
"output_tokens",
|
||||
"total_tokens",
|
||||
"cache_read_tokens",
|
||||
"cache_write_tokens",
|
||||
"reported_tokens",
|
||||
"estimated_tokens",
|
||||
"generation_ms",
|
||||
"measured_output_tokens",
|
||||
"ttft_ms",
|
||||
"timed_requests",
|
||||
)
|
||||
),
|
||||
_clean_status_code(call.error_status_code),
|
||||
_clean_error_kind(call.error_kind),
|
||||
)
|
||||
with self._lock:
|
||||
connection = self._connect()
|
||||
connection.execute(
|
||||
"""
|
||||
INSERT INTO llm_calls (
|
||||
started_at_ms, duration_ms, provider, model, source, stream,
|
||||
finish_reason, input_tokens, output_tokens, total_tokens,
|
||||
cache_read_tokens, cache_write_tokens, reported_tokens,
|
||||
estimated_tokens, generation_ms, measured_output_tokens,
|
||||
ttft_ms, timed_requests, error_status_code, error_kind
|
||||
) VALUES (
|
||||
?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?
|
||||
)
|
||||
""",
|
||||
values,
|
||||
)
|
||||
self._write_version += 1
|
||||
self._cached_payload_key = None
|
||||
self._cached_payload = None
|
||||
self._prune_if_due(connection)
|
||||
|
||||
def _prune_if_due(self, connection: sqlite3.Connection) -> None:
|
||||
utc_day = int(time.time() // 86_400)
|
||||
self._writes_since_size_prune += 1
|
||||
prune_age = self._last_prune_utc_day != utc_day
|
||||
prune_size = self._writes_since_size_prune >= 1_024
|
||||
if not prune_age and not prune_size:
|
||||
return
|
||||
if prune_age:
|
||||
cutoff_ms = int(
|
||||
(datetime.now(timezone.utc) - timedelta(days=MAX_DAYS_RETAINED)).timestamp()
|
||||
* 1000
|
||||
)
|
||||
connection.execute("DELETE FROM llm_calls WHERE started_at_ms < ?", (cutoff_ms,))
|
||||
connection.execute(
|
||||
"""
|
||||
DELETE FROM llm_calls
|
||||
WHERE id <= COALESCE((
|
||||
SELECT id FROM llm_calls ORDER BY id DESC LIMIT 1 OFFSET ?
|
||||
), -1)
|
||||
""",
|
||||
(MAX_CALLS_RETAINED,),
|
||||
)
|
||||
self._last_prune_utc_day = utc_day
|
||||
self._writes_since_size_prune = 0
|
||||
|
||||
def count(self) -> int:
|
||||
with self._lock:
|
||||
row = self._connect().execute("SELECT COUNT(*) AS count FROM llm_calls").fetchone()
|
||||
return int(row["count"] if row is not None else 0)
|
||||
|
||||
def _aggregate(
|
||||
self,
|
||||
*,
|
||||
connection: sqlite3.Connection,
|
||||
start_ms: int | None,
|
||||
end_ms: int,
|
||||
group_by: tuple[str, ...] = (),
|
||||
limit: int | None = None,
|
||||
) -> list[sqlite3.Row]:
|
||||
selected = f"{', '.join(group_by)}, " if group_by else ""
|
||||
where = "started_at_ms < ?"
|
||||
params: list[object] = [end_ms]
|
||||
if start_ms is not None:
|
||||
where = "started_at_ms >= ? AND started_at_ms < ?"
|
||||
params = [start_ms, end_ms]
|
||||
query = f"SELECT {selected}{_AGGREGATE_SQL} FROM llm_calls WHERE {where}"
|
||||
if group_by:
|
||||
query += f" GROUP BY {', '.join(group_by)} ORDER BY total_tokens DESC"
|
||||
if limit is not None:
|
||||
query += " LIMIT ?"
|
||||
params.append(limit)
|
||||
return list(connection.execute(query, params).fetchall())
|
||||
|
||||
def _daily_rows(
|
||||
self,
|
||||
*,
|
||||
connection: sqlite3.Connection,
|
||||
start_ms: int,
|
||||
end_ms: int,
|
||||
timezone_name: str,
|
||||
) -> list[dict[str, Any]]:
|
||||
query = f"""
|
||||
SELECT llm_usage_local_day(started_at_ms, ?) AS date, source,
|
||||
{_AGGREGATE_SQL}
|
||||
FROM llm_calls
|
||||
WHERE started_at_ms >= ? AND started_at_ms < ?
|
||||
GROUP BY date, source
|
||||
ORDER BY date, source
|
||||
"""
|
||||
rows = connection.execute(
|
||||
query,
|
||||
(timezone_name, start_ms, end_ms),
|
||||
).fetchall()
|
||||
by_date: dict[str, dict[str, Any]] = {}
|
||||
for row in rows:
|
||||
day = cast(str | None, row["date"])
|
||||
if day is None:
|
||||
continue
|
||||
values = _as_int_row(row)
|
||||
aggregate = by_date.setdefault(
|
||||
day,
|
||||
{"date": day, **_empty_totals(), "sources": {}},
|
||||
)
|
||||
for key, value in values.items():
|
||||
aggregate[key] += value
|
||||
aggregate["sources"][str(row["source"])] = values
|
||||
return list(by_date.values())
|
||||
|
||||
@staticmethod
|
||||
def _midnight_ms(value: date, zone: timezone | ZoneInfo) -> int:
|
||||
return int(datetime.combine(value, datetime.min.time(), tzinfo=zone).timestamp() * 1000)
|
||||
|
||||
def usage_payload(
|
||||
self,
|
||||
*,
|
||||
days: int = 371,
|
||||
timezone_name: str | None = None,
|
||||
now: datetime | None = None,
|
||||
) -> dict[str, Any]:
|
||||
zone = _zone(timezone_name)
|
||||
current = now or datetime.now(timezone.utc)
|
||||
if current.tzinfo is None:
|
||||
current = current.replace(tzinfo=timezone.utc)
|
||||
today = current.astimezone(zone).date()
|
||||
safe_days = max(1, days)
|
||||
zone_name = getattr(zone, "key", "UTC")
|
||||
|
||||
with self._lock:
|
||||
data_version_row = self._connect().execute("PRAGMA data_version").fetchone()
|
||||
data_version = int(data_version_row[0]) if data_version_row is not None else 0
|
||||
write_version = self._write_version
|
||||
cache_key = (
|
||||
safe_days,
|
||||
zone_name,
|
||||
today.isoformat(),
|
||||
write_version,
|
||||
data_version,
|
||||
)
|
||||
if self._cached_payload_key == cache_key and self._cached_payload is not None:
|
||||
return deepcopy(self._cached_payload)
|
||||
|
||||
connection = self._read_connection()
|
||||
try:
|
||||
connection.execute("BEGIN")
|
||||
end_ms = self._midnight_ms(today + timedelta(days=1), zone)
|
||||
retained_start = today - timedelta(days=MAX_DAYS_RETAINED - 1)
|
||||
retained_start_ms = self._midnight_ms(retained_start, zone)
|
||||
daily = self._daily_rows(
|
||||
connection=connection,
|
||||
start_ms=retained_start_ms,
|
||||
end_ms=end_ms,
|
||||
timezone_name=zone_name,
|
||||
)
|
||||
|
||||
requested_start = today - timedelta(days=safe_days - 1)
|
||||
visible_days = [row for row in daily if row["date"] >= requested_start.isoformat()]
|
||||
last_30_start_ms = self._midnight_ms(today - timedelta(days=29), zone)
|
||||
|
||||
last_30_date = (today - timedelta(days=29)).isoformat()
|
||||
last_365_date = (today - timedelta(days=364)).isoformat()
|
||||
all_totals = _sum_rows(daily)
|
||||
totals_30 = _sum_rows(row for row in daily if row["date"] >= last_30_date)
|
||||
totals_365 = _sum_rows(row for row in daily if row["date"] >= last_365_date)
|
||||
|
||||
provider_rows = self._aggregate(
|
||||
connection=connection,
|
||||
start_ms=last_30_start_ms,
|
||||
end_ms=end_ms,
|
||||
group_by=("provider", "model"),
|
||||
limit=50,
|
||||
)
|
||||
providers_30d = [
|
||||
{
|
||||
"provider": str(row["provider"]),
|
||||
"model": str(row["model"]),
|
||||
**_as_int_row(row),
|
||||
}
|
||||
for row in provider_rows
|
||||
]
|
||||
|
||||
active_dates = {
|
||||
date.fromisoformat(row["date"]) for row in daily if row["total_tokens"] > 0
|
||||
}
|
||||
current_streak = 0
|
||||
cursor = today
|
||||
while cursor in active_dates:
|
||||
current_streak += 1
|
||||
cursor -= timedelta(days=1)
|
||||
longest_streak = 0
|
||||
running_streak = 0
|
||||
previous: date | None = None
|
||||
for cursor in sorted(active_dates):
|
||||
running_streak = running_streak + 1 if previous == cursor - timedelta(days=1) else 1
|
||||
longest_streak = max(longest_streak, running_streak)
|
||||
previous = cursor
|
||||
|
||||
latest = (
|
||||
connection
|
||||
.execute("SELECT MAX(started_at_ms) AS updated_at_ms FROM llm_calls")
|
||||
.fetchone()
|
||||
)
|
||||
updated_at_ms = int(latest["updated_at_ms"] or 0) if latest is not None else 0
|
||||
denominator = totals_30["cache_read_observed_input_tokens"]
|
||||
payload = {
|
||||
"days": visible_days,
|
||||
"total_tokens": all_totals["total_tokens"],
|
||||
"total_tokens_30d": totals_30["total_tokens"],
|
||||
"total_tokens_365d": totals_365["total_tokens"],
|
||||
"reported_tokens_30d": totals_30["reported_tokens"],
|
||||
"estimated_tokens_30d": totals_30["estimated_tokens"],
|
||||
"cache_read_tokens_30d": totals_30["cache_read_tokens"],
|
||||
"cache_read_observed_input_tokens_30d": denominator,
|
||||
"cache_read_rate_30d": (
|
||||
totals_30["cache_read_tokens"] / denominator if denominator else None
|
||||
),
|
||||
"peak_day_tokens": max(
|
||||
(int(row["total_tokens"]) for row in daily),
|
||||
default=0,
|
||||
),
|
||||
"current_streak_days": current_streak,
|
||||
"longest_streak_days": longest_streak,
|
||||
"active_days_30d": sum(
|
||||
1
|
||||
for row in daily
|
||||
if row["date"] >= last_30_date and row["total_tokens"] > 0
|
||||
),
|
||||
"requests_30d": totals_30["requests"],
|
||||
"failed_requests_30d": totals_30["failed_requests"],
|
||||
"providers_30d": providers_30d,
|
||||
"updated_at": (
|
||||
datetime.fromtimestamp(updated_at_ms / 1000, timezone.utc)
|
||||
.isoformat()
|
||||
.replace("+00:00", "Z")
|
||||
if updated_at_ms
|
||||
else None
|
||||
),
|
||||
}
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
with self._lock:
|
||||
latest_data_version_row = self._connect().execute("PRAGMA data_version").fetchone()
|
||||
latest_data_version = (
|
||||
int(latest_data_version_row[0])
|
||||
if latest_data_version_row is not None
|
||||
else 0
|
||||
)
|
||||
if self._write_version == write_version and latest_data_version == data_version:
|
||||
self._cached_payload_key = cache_key
|
||||
self._cached_payload = payload
|
||||
return deepcopy(payload)
|
||||
|
||||
def recent_calls(self, *, limit: int = 100) -> list[dict[str, Any]]:
|
||||
"""Return bounded metadata rows for diagnostics; never returns content."""
|
||||
safe_limit = min(max(1, limit), 1_000)
|
||||
with self._lock:
|
||||
rows = (
|
||||
self._connect()
|
||||
.execute(
|
||||
"""
|
||||
SELECT * FROM llm_calls ORDER BY started_at_ms DESC, id DESC LIMIT ?
|
||||
""",
|
||||
(safe_limit,),
|
||||
)
|
||||
.fetchall()
|
||||
)
|
||||
return [dict(row) for row in rows]
|
||||
|
||||
def record_many(self, calls: Iterable[LLMCallRecord]) -> None:
|
||||
for call in calls:
|
||||
self.record(call)
|
||||
+143
-7
@@ -6,6 +6,7 @@ import asyncio
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from abc import ABC, abstractmethod
|
||||
from collections.abc import Awaitable, Callable
|
||||
from contextlib import suppress
|
||||
@@ -13,19 +14,23 @@ from copy import deepcopy
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from email.utils import parsedate_to_datetime
|
||||
from typing import Any, Literal, cast
|
||||
from typing import TYPE_CHECKING, Any, Literal, cast
|
||||
|
||||
import json_repair
|
||||
from loguru import logger
|
||||
|
||||
from nanobot.utils.helpers import sanitize_surrogates_deep
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from nanobot.llm_usage.models import LLMCallRecord
|
||||
|
||||
STREAM_IDLE_TIMEOUT_ENV = "NANOBOT_STREAM_IDLE_TIMEOUT_S"
|
||||
DEFAULT_STREAM_IDLE_TIMEOUT_S = 90.0
|
||||
MAX_STREAM_IDLE_TIMEOUT_S = 3600.0
|
||||
RETRY_AFTER_BUFFER = 1
|
||||
|
||||
RetryEventCallback = Callable[[str], Awaitable[None]]
|
||||
LLMCallObserver = Callable[["LLMCallRecord"], None]
|
||||
|
||||
|
||||
def resolve_stream_idle_timeout_s(
|
||||
@@ -682,6 +687,95 @@ class LLMProvider(ABC):
|
||||
self.api_base = api_base
|
||||
self.provider_name = provider_name
|
||||
self.generation: GenerationSettings = GenerationSettings()
|
||||
self._llm_call_observer: LLMCallObserver | None = None
|
||||
|
||||
def set_llm_call_observer(self, observer: LLMCallObserver | None) -> None:
|
||||
"""Attach a fail-open observer for each physical retry-managed call."""
|
||||
self._llm_call_observer = observer
|
||||
|
||||
def _usage_for_call(
|
||||
self,
|
||||
response: LLMResponse,
|
||||
kwargs: dict[str, Any],
|
||||
) -> LLMUsage | None:
|
||||
usage = response.usage
|
||||
if usage is None or usage.total_tokens == 0:
|
||||
if response.finish_reason in {"error", "cancelled"}:
|
||||
return None
|
||||
messages = kwargs.get("messages")
|
||||
if not isinstance(messages, list):
|
||||
return usage
|
||||
tools_value = kwargs.get("tools")
|
||||
tools = cast(list[dict[str, Any]], tools_value) if isinstance(tools_value, list) else None
|
||||
model_value = kwargs.get("model")
|
||||
model = model_value if isinstance(model_value, str) else self.get_default_model()
|
||||
try:
|
||||
from nanobot.utils.helpers import (
|
||||
build_assistant_message,
|
||||
estimate_message_tokens,
|
||||
estimate_prompt_tokens_chain,
|
||||
)
|
||||
|
||||
input_tokens, _ = estimate_prompt_tokens_chain(
|
||||
self,
|
||||
model,
|
||||
cast(list[dict[str, Any]], messages),
|
||||
tools,
|
||||
)
|
||||
assistant_message = build_assistant_message(
|
||||
response.content or "",
|
||||
tool_calls=[call.to_openai_tool_call() for call in response.tool_calls],
|
||||
reasoning_content=response.reasoning_content,
|
||||
thinking_blocks=response.thinking_blocks,
|
||||
)
|
||||
usage = LLMUsage.estimated(
|
||||
input_tokens=max(0, input_tokens),
|
||||
output_tokens=max(0, estimate_message_tokens(assistant_message)),
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("failed to estimate usage for {}", self.provider_name)
|
||||
return usage
|
||||
return usage.with_timing(
|
||||
generation_ms=response.generation_ms,
|
||||
ttft_ms=response.ttft_ms,
|
||||
)
|
||||
|
||||
def _observe_llm_call(
|
||||
self,
|
||||
response: LLMResponse,
|
||||
kwargs: dict[str, Any],
|
||||
*,
|
||||
started_at_ms: int,
|
||||
started_at_ns: int,
|
||||
stream: bool,
|
||||
) -> LLMResponse:
|
||||
observer = self._llm_call_observer
|
||||
if observer is None:
|
||||
return response
|
||||
usage = self._usage_for_call(response, kwargs)
|
||||
if usage is not None:
|
||||
response.usage = usage
|
||||
model_value = kwargs.get("model")
|
||||
model = model_value if isinstance(model_value, str) and model_value else self.get_default_model()
|
||||
try:
|
||||
from nanobot.llm_usage.context import current_llm_usage_source
|
||||
from nanobot.llm_usage.models import LLMCallRecord
|
||||
|
||||
observer(LLMCallRecord(
|
||||
started_at_ms=started_at_ms,
|
||||
duration_ms=max(0, (time.monotonic_ns() - started_at_ns) // 1_000_000),
|
||||
provider=self.provider_name,
|
||||
model=model,
|
||||
source=current_llm_usage_source(),
|
||||
stream=stream,
|
||||
finish_reason=response.finish_reason,
|
||||
usage=usage,
|
||||
error_status_code=response.error_status_code,
|
||||
error_kind=response.error_kind,
|
||||
))
|
||||
except Exception:
|
||||
logger.exception("LLM call observer failed for {}", self.provider_name)
|
||||
return response
|
||||
|
||||
def can_resume_conversation_state(
|
||||
self,
|
||||
@@ -1068,18 +1162,39 @@ class LLMProvider(ABC):
|
||||
|
||||
async def _safe_chat(self, **kwargs: Any) -> LLMResponse:
|
||||
"""Call chat() and convert unexpected exceptions to error responses."""
|
||||
started_at_ms = time.time_ns() // 1_000_000
|
||||
started_at_ns = time.monotonic_ns()
|
||||
try:
|
||||
provider_context = kwargs.pop("provider_context", None)
|
||||
if isinstance(provider_context, ProviderCallContext):
|
||||
return await self.chat_with_context(
|
||||
response = await self.chat_with_context(
|
||||
provider_context=provider_context,
|
||||
**kwargs,
|
||||
)
|
||||
return await self.chat(**kwargs)
|
||||
else:
|
||||
response = await self.chat(**kwargs)
|
||||
except asyncio.CancelledError:
|
||||
self._observe_llm_call(
|
||||
LLMResponse(
|
||||
content=None,
|
||||
finish_reason="cancelled",
|
||||
error_kind="cancelled",
|
||||
),
|
||||
kwargs,
|
||||
started_at_ms=started_at_ms,
|
||||
started_at_ns=started_at_ns,
|
||||
stream=False,
|
||||
)
|
||||
raise
|
||||
except Exception as exc:
|
||||
return LLMResponse(content=f"Error calling LLM: {exc}", finish_reason="error")
|
||||
response = LLMResponse(content=f"Error calling LLM: {exc}", finish_reason="error")
|
||||
return self._observe_llm_call(
|
||||
response,
|
||||
kwargs,
|
||||
started_at_ms=started_at_ms,
|
||||
started_at_ns=started_at_ns,
|
||||
stream=False,
|
||||
)
|
||||
|
||||
async def chat_stream(
|
||||
self,
|
||||
@@ -1142,18 +1257,39 @@ class LLMProvider(ABC):
|
||||
|
||||
async def _safe_chat_stream(self, **kwargs: Any) -> LLMResponse:
|
||||
"""Call chat_stream() and convert unexpected exceptions to error responses."""
|
||||
started_at_ms = time.time_ns() // 1_000_000
|
||||
started_at_ns = time.monotonic_ns()
|
||||
try:
|
||||
provider_context = kwargs.pop("provider_context", None)
|
||||
if isinstance(provider_context, ProviderCallContext):
|
||||
return await self.chat_stream_with_context(
|
||||
response = await self.chat_stream_with_context(
|
||||
provider_context=provider_context,
|
||||
**kwargs,
|
||||
)
|
||||
return await self.chat_stream(**kwargs)
|
||||
else:
|
||||
response = await self.chat_stream(**kwargs)
|
||||
except asyncio.CancelledError:
|
||||
self._observe_llm_call(
|
||||
LLMResponse(
|
||||
content=None,
|
||||
finish_reason="cancelled",
|
||||
error_kind="cancelled",
|
||||
),
|
||||
kwargs,
|
||||
started_at_ms=started_at_ms,
|
||||
started_at_ns=started_at_ns,
|
||||
stream=True,
|
||||
)
|
||||
raise
|
||||
except Exception as exc:
|
||||
return LLMResponse(content=f"Error calling LLM: {exc}", finish_reason="error")
|
||||
response = LLMResponse(content=f"Error calling LLM: {exc}", finish_reason="error")
|
||||
return self._observe_llm_call(
|
||||
response,
|
||||
kwargs,
|
||||
started_at_ms=started_at_ms,
|
||||
started_at_ns=started_at_ns,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
async def chat_stream_with_retry(
|
||||
self,
|
||||
|
||||
@@ -13,6 +13,7 @@ from loguru import logger
|
||||
|
||||
from nanobot.providers.base import (
|
||||
GenerationSettings,
|
||||
LLMCallObserver,
|
||||
LLMProvider,
|
||||
LLMResponse,
|
||||
ProviderCallContext,
|
||||
@@ -151,6 +152,11 @@ class FallbackProvider(LLMProvider):
|
||||
"""Attach a process-level observer without changing request call signatures."""
|
||||
self._fallback_model_observer = observer
|
||||
|
||||
def set_llm_call_observer(self, observer: LLMCallObserver | None) -> None:
|
||||
"""Attach usage recording to the primary and future fallback leaves."""
|
||||
super().set_llm_call_observer(observer)
|
||||
self._primary.set_llm_call_observer(observer)
|
||||
|
||||
@property
|
||||
def supports_progress_deltas(self) -> bool:
|
||||
return bool(getattr(self._primary, "supports_progress_deltas", False))
|
||||
@@ -506,6 +512,7 @@ class FallbackProvider(LLMProvider):
|
||||
)
|
||||
try:
|
||||
fallback_provider = self._provider_factory(fallback)
|
||||
fallback_provider.set_llm_call_observer(self._llm_call_observer)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Failed to create provider for fallback '{}': {}", fallback_model, exc
|
||||
|
||||
@@ -37,6 +37,7 @@ from nanobot.bus.runtime_events import (
|
||||
TurnRuntimeAdmitted,
|
||||
UserInputAccepted,
|
||||
)
|
||||
from nanobot.llm_usage.context import llm_usage_source
|
||||
from nanobot.providers.base import LLMProvider, LLMUsage
|
||||
from nanobot.providers.fallback_provider import FallbackModelObserver
|
||||
from nanobot.runtime_context import public_history_message
|
||||
@@ -208,24 +209,25 @@ async def maybe_generate_webui_title(
|
||||
prompt += f"\nAssistant: {truncate_text(assistant_text, 1_000)}"
|
||||
|
||||
try:
|
||||
response = await provider.chat_with_retry(
|
||||
[
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"You write short, neutral chat titles. "
|
||||
"Return only the title text."
|
||||
),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
tools=None,
|
||||
model=model,
|
||||
max_tokens=TITLE_GENERATION_MAX_TOKENS,
|
||||
temperature=0.2,
|
||||
reasoning_effort=TITLE_GENERATION_REASONING_EFFORT,
|
||||
retry_mode="standard",
|
||||
)
|
||||
with llm_usage_source("system"):
|
||||
response = await provider.chat_with_retry(
|
||||
[
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"You write short, neutral chat titles. "
|
||||
"Return only the title text."
|
||||
),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
tools=None,
|
||||
model=model,
|
||||
max_tokens=TITLE_GENERATION_MAX_TOKENS,
|
||||
temperature=0.2,
|
||||
reasoning_effort=TITLE_GENERATION_REASONING_EFFORT,
|
||||
retry_mode="standard",
|
||||
)
|
||||
except Exception:
|
||||
logger.debug("Failed to generate webui session title for {}", session_key, exc_info=True)
|
||||
return False
|
||||
|
||||
@@ -284,9 +284,9 @@ class WebUISettingsRouter:
|
||||
if not self._authorized(request):
|
||||
return self._unauthorized()
|
||||
if route == ("root", "settings"):
|
||||
return self._handle_settings()
|
||||
return await asyncio.to_thread(self._handle_settings)
|
||||
if route == ("root", "usage"):
|
||||
return self._handle_settings_usage()
|
||||
return await asyncio.to_thread(self._handle_settings_usage)
|
||||
|
||||
domain, action = route
|
||||
domain_request = self._domain_request(
|
||||
|
||||
@@ -20,6 +20,7 @@ from nanobot.channels.contracts import (
|
||||
channel_update_instance_config,
|
||||
)
|
||||
from nanobot.config.schema import Config
|
||||
from nanobot.llm_usage import llm_usage_payload
|
||||
from nanobot.optional_features import OptionalFeatureError, with_channel_runtime_status
|
||||
from nanobot.security.workspace_access import workspace_sandbox_status
|
||||
from nanobot.webui.settings_capabilities import network_safety_payload
|
||||
@@ -31,7 +32,6 @@ from nanobot.webui.settings_contracts import (
|
||||
query_first,
|
||||
query_first_alias,
|
||||
)
|
||||
from nanobot.webui.token_usage import token_usage_payload
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from nanobot.webui.settings_services import WebUISettingsServices
|
||||
@@ -121,7 +121,7 @@ def system_settings_payload(
|
||||
},
|
||||
"unified_session": defaults.unified_session,
|
||||
},
|
||||
"usage": token_usage_payload(timezone_name=defaults.timezone),
|
||||
"usage": llm_usage_payload(timezone_name=defaults.timezone),
|
||||
"advanced": {
|
||||
"restrict_to_workspace": config.tools.restrict_to_workspace,
|
||||
"workspace_sandbox": sandbox_status.as_dict(),
|
||||
@@ -139,7 +139,7 @@ def system_settings_payload(
|
||||
|
||||
def settings_usage_payload(config: Config) -> dict[str, Any]:
|
||||
"""Return the lightweight token usage slice for Overview refreshes."""
|
||||
return token_usage_payload(timezone_name=config.agents.defaults.timezone)
|
||||
return llm_usage_payload(timezone_name=config.agents.defaults.timezone)
|
||||
|
||||
|
||||
def update_agent_system_settings(config: Config, query: QueryParams) -> tuple[bool, bool]:
|
||||
|
||||
@@ -1,392 +0,0 @@
|
||||
"""Workspace-scoped token usage telemetry for WebUI overview surfaces."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any, Mapping, cast
|
||||
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
|
||||
|
||||
from loguru import logger
|
||||
|
||||
from nanobot.agent.hook import AgentHook, AgentHookContext
|
||||
from nanobot.config.paths import get_webui_dir
|
||||
from nanobot.providers.base import LLMUsage
|
||||
|
||||
TOKEN_USAGE_SCHEMA_VERSION = 2
|
||||
_MAX_STATE_FILE_BYTES = 512 * 1024
|
||||
_MAX_DAYS_RETAINED = 400
|
||||
_USAGE_KEYS = (
|
||||
"input_tokens",
|
||||
"output_tokens",
|
||||
"cache_read_tokens",
|
||||
"cache_write_tokens",
|
||||
"cache_read_observed_input_tokens",
|
||||
"cache_write_observed_input_tokens",
|
||||
"total_tokens",
|
||||
"reported_tokens",
|
||||
"estimated_tokens",
|
||||
)
|
||||
_REQUEST_KEYS = ("requests", "reported_requests", "estimated_requests")
|
||||
_SOURCE_KEYS = ("user", "api", "cron", "dream", "system")
|
||||
_WRITE_LOCK = threading.Lock()
|
||||
|
||||
|
||||
def token_usage_state_path() -> Path:
|
||||
return get_webui_dir() / "token-usage.json"
|
||||
|
||||
|
||||
def default_token_usage_state() -> dict[str, Any]:
|
||||
return {
|
||||
"schema_version": TOKEN_USAGE_SCHEMA_VERSION,
|
||||
"days": {},
|
||||
"updated_at": None,
|
||||
}
|
||||
|
||||
|
||||
def _utc_now_iso() -> str:
|
||||
return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
|
||||
|
||||
|
||||
def _zone(timezone_name: str | None) -> timezone | ZoneInfo:
|
||||
if not timezone_name:
|
||||
return timezone.utc
|
||||
try:
|
||||
return ZoneInfo(timezone_name)
|
||||
except ZoneInfoNotFoundError:
|
||||
return timezone.utc
|
||||
|
||||
|
||||
def _local_day(now: datetime | None = None, *, timezone_name: str | None = None) -> str:
|
||||
dt = now or datetime.now(timezone.utc)
|
||||
if dt.tzinfo is None:
|
||||
dt = dt.replace(tzinfo=timezone.utc)
|
||||
return dt.astimezone(_zone(timezone_name)).date().isoformat()
|
||||
|
||||
|
||||
def _clean_int(value: Any) -> int:
|
||||
try:
|
||||
return max(0, int(value or 0))
|
||||
except (TypeError, ValueError):
|
||||
return 0
|
||||
|
||||
|
||||
def _clean_source(value: str | None) -> str:
|
||||
return value if value in _SOURCE_KEYS else "system"
|
||||
|
||||
|
||||
def _source_from_session_key(session_key: str | None) -> str:
|
||||
key = session_key or ""
|
||||
if key.startswith("dream:"):
|
||||
return "dream"
|
||||
if key == "heartbeat" or key.startswith("cron:"):
|
||||
return "cron"
|
||||
if key.startswith("api:"):
|
||||
return "api"
|
||||
if key.startswith("system:"):
|
||||
return "system"
|
||||
return "user"
|
||||
|
||||
|
||||
def _normalize_usage(raw: LLMUsage | None) -> dict[str, int]:
|
||||
if raw is None:
|
||||
return {}
|
||||
usage = {
|
||||
"input_tokens": raw.input_tokens,
|
||||
"output_tokens": raw.output_tokens,
|
||||
"cache_read_tokens": raw.cache_read_tokens or 0,
|
||||
"cache_write_tokens": raw.cache_write_tokens or 0,
|
||||
"cache_read_observed_input_tokens": (
|
||||
raw.input_tokens if raw.cache_read_tokens is not None else 0
|
||||
),
|
||||
"cache_write_observed_input_tokens": (
|
||||
raw.input_tokens if raw.cache_write_tokens is not None else 0
|
||||
),
|
||||
"total_tokens": raw.total_tokens,
|
||||
"reported_tokens": raw.reported_tokens,
|
||||
"estimated_tokens": raw.estimated_tokens,
|
||||
}
|
||||
return usage if usage["total_tokens"] > 0 else {}
|
||||
|
||||
|
||||
def _normalize_usage_row(row: dict[str, Any]) -> dict[str, int]:
|
||||
cleaned = {key: _clean_int(row.get(key)) for key in _USAGE_KEYS}
|
||||
if cleaned["total_tokens"] <= 0:
|
||||
cleaned["total_tokens"] = cleaned["input_tokens"] + cleaned["output_tokens"]
|
||||
if cleaned["reported_tokens"] <= 0 and cleaned["estimated_tokens"] <= 0:
|
||||
cleaned["reported_tokens"] = cleaned["total_tokens"]
|
||||
requests = {key: _clean_int(row.get(key)) for key in _REQUEST_KEYS}
|
||||
if (
|
||||
requests["requests"] > 0
|
||||
and requests["reported_requests"] <= 0
|
||||
and requests["estimated_requests"] <= 0
|
||||
):
|
||||
if cleaned["estimated_tokens"] > 0 and cleaned["reported_tokens"] <= 0:
|
||||
requests["estimated_requests"] = requests["requests"]
|
||||
else:
|
||||
requests["reported_requests"] = requests["requests"]
|
||||
return {**cleaned, **requests}
|
||||
|
||||
|
||||
def _normalize_sources(raw: Any, fallback: dict[str, int]) -> dict[str, dict[str, int]]:
|
||||
sources: dict[str, dict[str, int]] = {}
|
||||
if isinstance(raw, dict):
|
||||
for source, row_value in cast(dict[Any, Any], raw).items():
|
||||
if not isinstance(row_value, dict):
|
||||
continue
|
||||
row = cast(dict[str, Any], row_value)
|
||||
normalized = _normalize_usage_row(row)
|
||||
if normalized["total_tokens"] <= 0 and normalized["requests"] <= 0:
|
||||
continue
|
||||
source_key = _clean_source(str(source))
|
||||
current = sources.get(source_key)
|
||||
if current is None:
|
||||
sources[source_key] = normalized
|
||||
else:
|
||||
for key in (*_USAGE_KEYS, *_REQUEST_KEYS):
|
||||
current[key] = _clean_int(current.get(key)) + normalized[key]
|
||||
if not sources and (fallback["total_tokens"] > 0 or fallback["requests"] > 0):
|
||||
sources["user"] = {key: fallback[key] for key in (*_USAGE_KEYS, *_REQUEST_KEYS)}
|
||||
return sources
|
||||
|
||||
|
||||
def normalize_token_usage_state(raw: Any) -> dict[str, Any]:
|
||||
state = default_token_usage_state()
|
||||
if not isinstance(raw, dict):
|
||||
return state
|
||||
raw = cast(dict[str, Any], raw)
|
||||
if raw.get("schema_version") != TOKEN_USAGE_SCHEMA_VERSION:
|
||||
return state
|
||||
days_raw = raw.get("days")
|
||||
if not isinstance(days_raw, dict):
|
||||
return state
|
||||
|
||||
days: dict[str, dict[str, Any]] = {}
|
||||
for date, row_value in sorted(cast(dict[Any, Any], days_raw).items())[-_MAX_DAYS_RETAINED:]:
|
||||
if not isinstance(date, str) or len(date) != 10 or not isinstance(row_value, dict):
|
||||
continue
|
||||
row = cast(dict[str, Any], row_value)
|
||||
try:
|
||||
datetime.fromisoformat(date)
|
||||
except ValueError:
|
||||
# A hand-edited or foreign day key that is not a real date would
|
||||
# otherwise reach token_usage_payload's date parsing and fail every
|
||||
# settings request; drop it like any other malformed row.
|
||||
continue
|
||||
normalized = _normalize_usage_row(row)
|
||||
if normalized["total_tokens"] <= 0 and normalized["requests"] <= 0:
|
||||
continue
|
||||
days[date] = {
|
||||
"date": date,
|
||||
**normalized,
|
||||
"sources": _normalize_sources(row.get("sources"), normalized),
|
||||
}
|
||||
|
||||
state["days"] = days
|
||||
updated_at = raw.get("updated_at")
|
||||
state["updated_at"] = updated_at if isinstance(updated_at, str) else None
|
||||
return state
|
||||
|
||||
|
||||
def read_token_usage_state() -> dict[str, Any]:
|
||||
path = token_usage_state_path()
|
||||
if not path.is_file():
|
||||
return default_token_usage_state()
|
||||
try:
|
||||
if path.stat().st_size > _MAX_STATE_FILE_BYTES:
|
||||
logger.warning("token usage state too large, ignoring: {}", path)
|
||||
return default_token_usage_state()
|
||||
with open(path, encoding="utf-8") as f:
|
||||
raw = json.load(f)
|
||||
except (OSError, json.JSONDecodeError) as e:
|
||||
logger.warning("read token usage state failed {}: {}", path, e)
|
||||
return default_token_usage_state()
|
||||
return normalize_token_usage_state(raw)
|
||||
|
||||
|
||||
def _encode_token_usage_state(state: dict[str, Any]) -> bytes:
|
||||
"""Encode the persisted state compactly, including its trailing newline."""
|
||||
payload = json.dumps(
|
||||
state,
|
||||
ensure_ascii=False,
|
||||
separators=(",", ":"),
|
||||
sort_keys=True,
|
||||
)
|
||||
return f"{payload}\n".encode("utf-8")
|
||||
|
||||
|
||||
def write_token_usage_state(raw: dict[str, Any]) -> dict[str, Any]:
|
||||
# Day-count retention is applied by normalization first. The byte budget
|
||||
# then trims only the oldest remaining days, preserving a contiguous suffix.
|
||||
state = normalize_token_usage_state(raw)
|
||||
state["updated_at"] = _utc_now_iso()
|
||||
days = cast(dict[str, dict[str, Any]], state["days"])
|
||||
encoded = _encode_token_usage_state(state)
|
||||
while len(encoded) > _MAX_STATE_FILE_BYTES and len(days) > 1:
|
||||
del days[min(days)]
|
||||
encoded = _encode_token_usage_state(state)
|
||||
if len(encoded) > _MAX_STATE_FILE_BYTES:
|
||||
raise ValueError("latest token usage day exceeds the state byte limit")
|
||||
|
||||
path = token_usage_state_path()
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
tmp = path.with_suffix(".json.tmp")
|
||||
with open(tmp, "wb") as f:
|
||||
f.write(encoded)
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
os.replace(tmp, path)
|
||||
try:
|
||||
dir_fd = os.open(path.parent, os.O_RDONLY)
|
||||
except OSError:
|
||||
return state
|
||||
try:
|
||||
os.fsync(dir_fd)
|
||||
finally:
|
||||
os.close(dir_fd)
|
||||
return state
|
||||
|
||||
|
||||
def record_token_usage(
|
||||
usage: LLMUsage | None,
|
||||
*,
|
||||
source: str = "user",
|
||||
timezone_name: str | None = None,
|
||||
now: datetime | None = None,
|
||||
) -> dict[str, Any]:
|
||||
normalized = _normalize_usage(usage)
|
||||
if not normalized:
|
||||
return read_token_usage_state()
|
||||
|
||||
with _WRITE_LOCK:
|
||||
state = read_token_usage_state()
|
||||
days_by_date = cast(dict[str, dict[str, Any]], state["days"])
|
||||
day = _local_day(now, timezone_name=timezone_name)
|
||||
row: dict[str, Any] = dict(days_by_date.get(day) or {"date": day, "requests": 0})
|
||||
for key in _USAGE_KEYS:
|
||||
row[key] = _clean_int(row.get(key)) + normalized.get(key, 0)
|
||||
row["requests"] = _clean_int(row.get("requests")) + 1
|
||||
if normalized.get("estimated_tokens", 0) > 0 and normalized.get("reported_tokens", 0) <= 0:
|
||||
row["estimated_requests"] = _clean_int(row.get("estimated_requests")) + 1
|
||||
else:
|
||||
row["reported_requests"] = _clean_int(row.get("reported_requests")) + 1
|
||||
|
||||
source_key = _clean_source(source)
|
||||
sources: dict[str, dict[str, Any]] = dict(
|
||||
cast(Mapping[str, dict[str, Any]], row.get("sources") or {})
|
||||
)
|
||||
source_row: dict[str, Any] = dict(sources.get(source_key) or {"requests": 0})
|
||||
for key in _USAGE_KEYS:
|
||||
source_row[key] = _clean_int(source_row.get(key)) + normalized.get(key, 0)
|
||||
source_row["requests"] = _clean_int(source_row.get("requests")) + 1
|
||||
if normalized.get("estimated_tokens", 0) > 0 and normalized.get("reported_tokens", 0) <= 0:
|
||||
source_row["estimated_requests"] = _clean_int(source_row.get("estimated_requests")) + 1
|
||||
else:
|
||||
source_row["reported_requests"] = _clean_int(source_row.get("reported_requests")) + 1
|
||||
sources[source_key] = source_row
|
||||
row["sources"] = sources
|
||||
|
||||
days_by_date[day] = row
|
||||
if len(days_by_date) > _MAX_DAYS_RETAINED:
|
||||
state["days"] = dict(sorted(days_by_date.items())[-_MAX_DAYS_RETAINED:])
|
||||
return write_token_usage_state(state)
|
||||
|
||||
|
||||
def record_response_token_usage(
|
||||
response: Any,
|
||||
*,
|
||||
source: str,
|
||||
timezone_name: str | None = None,
|
||||
) -> None:
|
||||
try:
|
||||
record_token_usage(
|
||||
getattr(response, "usage", None),
|
||||
source=source,
|
||||
timezone_name=timezone_name,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("failed to record {} token usage", source)
|
||||
|
||||
|
||||
def token_usage_payload(
|
||||
*,
|
||||
days: int = 371,
|
||||
timezone_name: str | None = None,
|
||||
now: datetime | None = None,
|
||||
) -> dict[str, Any]:
|
||||
state = read_token_usage_state()
|
||||
days_by_date = cast(dict[str, dict[str, Any]], state["days"])
|
||||
today = datetime.fromisoformat(_local_day(now, timezone_name=timezone_name)).date()
|
||||
start = today - timedelta(days=max(1, days) - 1)
|
||||
day_rows = [
|
||||
row
|
||||
for date, row in sorted(days_by_date.items())
|
||||
if start.isoformat() <= date <= today.isoformat()
|
||||
]
|
||||
last_30_start = today - timedelta(days=29)
|
||||
last_30 = [
|
||||
row
|
||||
for date, row in days_by_date.items()
|
||||
if last_30_start.isoformat() <= date <= today.isoformat()
|
||||
]
|
||||
last_365_start = today - timedelta(days=364)
|
||||
last_365 = [
|
||||
row
|
||||
for date, row in days_by_date.items()
|
||||
if last_365_start.isoformat() <= date <= today.isoformat()
|
||||
]
|
||||
active_dates = {
|
||||
datetime.fromisoformat(date).date()
|
||||
for date, row in days_by_date.items()
|
||||
if _clean_int(row.get("total_tokens")) > 0
|
||||
}
|
||||
current_streak = 0
|
||||
cursor = today
|
||||
while cursor in active_dates:
|
||||
current_streak += 1
|
||||
cursor -= timedelta(days=1)
|
||||
|
||||
longest_streak = 0
|
||||
running_streak = 0
|
||||
for cursor in sorted(active_dates):
|
||||
if cursor - timedelta(days=1) in active_dates:
|
||||
running_streak += 1
|
||||
else:
|
||||
running_streak = 1
|
||||
longest_streak = max(longest_streak, running_streak)
|
||||
|
||||
all_rows = list(days_by_date.values())
|
||||
return {
|
||||
"days": day_rows,
|
||||
"total_tokens": sum(_clean_int(row.get("total_tokens")) for row in all_rows),
|
||||
"total_tokens_30d": sum(_clean_int(row.get("total_tokens")) for row in last_30),
|
||||
"total_tokens_365d": sum(_clean_int(row.get("total_tokens")) for row in last_365),
|
||||
"peak_day_tokens": max([_clean_int(row.get("total_tokens")) for row in all_rows] or [0]),
|
||||
"current_streak_days": current_streak,
|
||||
"longest_streak_days": longest_streak,
|
||||
"active_days_30d": sum(1 for row in last_30 if _clean_int(row.get("total_tokens")) > 0),
|
||||
"requests_30d": sum(_clean_int(row.get("requests")) for row in last_30),
|
||||
"updated_at": state.get("updated_at"),
|
||||
}
|
||||
|
||||
|
||||
class TokenUsageHook(AgentHook):
|
||||
"""Persist provider-reported token usage without coupling it to chat messages."""
|
||||
|
||||
def __init__(self, *, timezone_name: str | None = None) -> None:
|
||||
super().__init__()
|
||||
self._timezone_name = timezone_name
|
||||
|
||||
async def after_iteration(self, context: AgentHookContext) -> None:
|
||||
try:
|
||||
record_token_usage(
|
||||
context.usage,
|
||||
source=_source_from_session_key(context.session_key),
|
||||
timezone_name=self._timezone_name,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("failed to record token usage")
|
||||
@@ -1,5 +1,6 @@
|
||||
"""Tests for SubagentManager."""
|
||||
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
@@ -10,6 +11,7 @@ from nanobot.agent.subagent import SubagentManager, SubagentStatus
|
||||
from nanobot.agent.tools.filesystem import FileToolsConfig
|
||||
from nanobot.bus.queue import MessageBus
|
||||
from nanobot.config.schema import ToolsConfig
|
||||
from nanobot.llm_usage.context import llm_usage_source
|
||||
from nanobot.providers.base import GenerationSettings, LLMProvider
|
||||
from nanobot.security.workspace_access import build_workspace_scope
|
||||
from nanobot.utils.llm_runtime import LLMRuntime
|
||||
@@ -198,3 +200,30 @@ async def test_subagent_forwards_fail_on_tool_error_to_runner(tmp_path):
|
||||
|
||||
spec = sm.runner.run.call_args.args[0]
|
||||
assert spec.fail_on_tool_error is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_spawned_subagent_inherits_llm_usage_source(tmp_path):
|
||||
provider = MagicMock(spec=LLMProvider)
|
||||
provider.get_default_model.return_value = "test"
|
||||
sm = SubagentManager(
|
||||
workspace=tmp_path,
|
||||
bus=MessageBus(),
|
||||
max_tool_result_chars=16_000,
|
||||
)
|
||||
sm.runner.run = AsyncMock(
|
||||
return_value=AgentRunResult(final_content="ok", messages=[], stop_reason="completed")
|
||||
)
|
||||
sm._announce_result = AsyncMock()
|
||||
|
||||
with llm_usage_source("cron"):
|
||||
await sm.spawn(
|
||||
"automation task",
|
||||
session_key="websocket:bound-automation",
|
||||
runtime=_runtime(provider),
|
||||
)
|
||||
tasks = list(sm._running_tasks.values())
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
spec = sm.runner.run.call_args.args[0]
|
||||
assert spec.llm_usage_source == "cron"
|
||||
|
||||
@@ -0,0 +1,19 @@
|
||||
from nanobot.llm_usage.context import source_from_request
|
||||
|
||||
|
||||
def test_automation_metadata_overrides_user_session_source() -> None:
|
||||
assert source_from_request(
|
||||
"websocket:ordinary-session",
|
||||
channel="websocket",
|
||||
metadata={"_cron_trigger": {"job_id": "job"}},
|
||||
) == "cron"
|
||||
assert source_from_request(
|
||||
"websocket:ordinary-session",
|
||||
channel="websocket",
|
||||
metadata={"_local_trigger": {"trigger_id": "trigger"}},
|
||||
) == "cron"
|
||||
|
||||
|
||||
def test_api_and_system_channels_have_explicit_sources() -> None:
|
||||
assert source_from_request("shared-session", channel="api", metadata={}) == "api"
|
||||
assert source_from_request("shared-session", channel="system", metadata={}) == "system"
|
||||
@@ -0,0 +1,275 @@
|
||||
"""Tests for the SQLite LLM usage store."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlite3
|
||||
import threading
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from nanobot.llm_usage.context import LLMUsageSource
|
||||
from nanobot.llm_usage.models import LLMCallRecord
|
||||
from nanobot.llm_usage.store import SCHEMA_VERSION, LLMUsageStore
|
||||
from nanobot.providers.base import LLMUsage
|
||||
|
||||
|
||||
def _timestamp(value: str) -> int:
|
||||
return int(datetime.fromisoformat(value).timestamp() * 1000)
|
||||
|
||||
|
||||
def _call(
|
||||
started_at: str,
|
||||
*,
|
||||
provider: str = "openai",
|
||||
model: str = "gpt-5",
|
||||
source: LLMUsageSource = "user",
|
||||
usage: LLMUsage | None = None,
|
||||
finish_reason: str = "stop",
|
||||
error_kind: str | None = None,
|
||||
) -> LLMCallRecord:
|
||||
return LLMCallRecord(
|
||||
started_at_ms=_timestamp(started_at),
|
||||
duration_ms=250,
|
||||
provider=provider,
|
||||
model=model,
|
||||
source=source,
|
||||
stream=True,
|
||||
finish_reason=finish_reason,
|
||||
usage=usage,
|
||||
error_status_code=429 if finish_reason == "error" else None,
|
||||
error_kind=error_kind or ("rate_limit" if finish_reason == "error" else None),
|
||||
)
|
||||
|
||||
|
||||
def test_store_keeps_only_content_free_call_metadata(tmp_path: Path) -> None:
|
||||
path = tmp_path / "llm_usage.sqlite3"
|
||||
store = LLMUsageStore(path)
|
||||
store.record(
|
||||
_call(
|
||||
"2026-06-03T00:00:00+00:00",
|
||||
usage=LLMUsage.reported(input_tokens=100, output_tokens=20),
|
||||
)
|
||||
)
|
||||
|
||||
row = store.recent_calls(limit=1)[0]
|
||||
assert row["provider"] == "openai"
|
||||
assert row["model"] == "gpt-5"
|
||||
assert row["total_tokens"] == 120
|
||||
assert not {
|
||||
"messages",
|
||||
"prompt",
|
||||
"content",
|
||||
"response",
|
||||
"tool_calls",
|
||||
"error_type",
|
||||
"error_code",
|
||||
} & set(row)
|
||||
|
||||
with sqlite3.connect(path) as connection:
|
||||
version = connection.execute("PRAGMA user_version").fetchone()[0]
|
||||
mode = connection.execute("PRAGMA journal_mode").fetchone()[0]
|
||||
assert version == SCHEMA_VERSION
|
||||
assert str(mode).lower() == "wal"
|
||||
|
||||
|
||||
def test_usage_payload_aggregates_cache_coverage_sources_and_failures(tmp_path: Path) -> None:
|
||||
store = LLMUsageStore(tmp_path / "llm_usage.sqlite3")
|
||||
store.record_many(
|
||||
[
|
||||
_call(
|
||||
"2026-06-02T23:30:00+00:00",
|
||||
usage=LLMUsage.reported(
|
||||
input_tokens=100,
|
||||
output_tokens=20,
|
||||
cache_read_tokens=40,
|
||||
cache_write_tokens=10,
|
||||
),
|
||||
),
|
||||
_call(
|
||||
"2026-06-03T01:00:00+00:00",
|
||||
source="api",
|
||||
usage=LLMUsage.reported(input_tokens=50, output_tokens=5),
|
||||
),
|
||||
_call(
|
||||
"2026-06-03T02:00:00+00:00",
|
||||
provider="anthropic",
|
||||
model="claude-sonnet-4",
|
||||
source="dream",
|
||||
usage=LLMUsage.estimated(input_tokens=30, output_tokens=10),
|
||||
),
|
||||
_call(
|
||||
"2026-06-03T03:00:00+00:00",
|
||||
provider="anthropic",
|
||||
model="claude-sonnet-4",
|
||||
source="system",
|
||||
finish_reason="error",
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
payload = store.usage_payload(
|
||||
timezone_name="Asia/Shanghai",
|
||||
now=datetime(2026, 6, 3, 8, tzinfo=timezone.utc),
|
||||
)
|
||||
|
||||
assert payload["total_tokens_30d"] == 215
|
||||
assert payload["reported_tokens_30d"] == 175
|
||||
assert payload["estimated_tokens_30d"] == 40
|
||||
assert payload["requests_30d"] == 4
|
||||
assert payload["failed_requests_30d"] == 1
|
||||
assert payload["cache_read_tokens_30d"] == 40
|
||||
assert payload["cache_read_observed_input_tokens_30d"] == 100
|
||||
assert payload["cache_read_rate_30d"] == 0.4
|
||||
|
||||
day = payload["days"][0]
|
||||
assert day["date"] == "2026-06-03"
|
||||
assert day["requests"] == 4
|
||||
assert day["reported_requests"] == 2
|
||||
assert day["estimated_requests"] == 1
|
||||
assert day["sources"]["api"]["cache_read_observed_input_tokens"] == 0
|
||||
assert day["sources"]["user"]["cache_read_observed_input_tokens"] == 100
|
||||
assert {(row["provider"], row["model"]) for row in payload["providers_30d"]} == {
|
||||
("openai", "gpt-5"),
|
||||
("anthropic", "claude-sonnet-4"),
|
||||
}
|
||||
|
||||
|
||||
def test_usage_payload_preserves_zero_cache_observation(tmp_path: Path) -> None:
|
||||
store = LLMUsageStore(tmp_path / "llm_usage.sqlite3")
|
||||
store.record(
|
||||
_call(
|
||||
"2026-06-03T00:00:00+00:00",
|
||||
usage=LLMUsage.reported(
|
||||
input_tokens=80,
|
||||
output_tokens=5,
|
||||
cache_read_tokens=0,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
payload = store.usage_payload(
|
||||
now=datetime(2026, 6, 3, 12, tzinfo=timezone.utc),
|
||||
)
|
||||
|
||||
assert payload["cache_read_tokens_30d"] == 0
|
||||
assert payload["cache_read_observed_input_tokens_30d"] == 80
|
||||
assert payload["cache_read_rate_30d"] == 0.0
|
||||
|
||||
|
||||
def test_recent_calls_is_bounded(tmp_path: Path) -> None:
|
||||
store = LLMUsageStore(tmp_path / "llm_usage.sqlite3")
|
||||
call = _call(
|
||||
"2026-06-03T00:00:00+00:00",
|
||||
usage=LLMUsage.reported(input_tokens=1, output_tokens=1),
|
||||
)
|
||||
store.record_many(call for _ in range(1_005))
|
||||
|
||||
assert len(store.recent_calls(limit=10_000)) == 1_000
|
||||
|
||||
|
||||
def test_cancelled_calls_are_failures_and_error_kind_is_coarse(tmp_path: Path) -> None:
|
||||
store = LLMUsageStore(tmp_path / "llm_usage.sqlite3")
|
||||
call = _call(
|
||||
"2026-06-03T00:00:00+00:00",
|
||||
finish_reason="cancelled",
|
||||
error_kind="provider payload: secret text",
|
||||
)
|
||||
store.record(call)
|
||||
|
||||
payload = store.usage_payload(
|
||||
now=datetime(2026, 6, 3, 12, tzinfo=timezone.utc),
|
||||
)
|
||||
|
||||
assert payload["failed_requests_30d"] == 1
|
||||
assert store.recent_calls(limit=1)[0]["error_kind"] == "other"
|
||||
|
||||
|
||||
def test_usage_payload_cache_is_isolated_and_invalidated_on_write(tmp_path: Path) -> None:
|
||||
store = LLMUsageStore(tmp_path / "llm_usage.sqlite3")
|
||||
first_call = _call(
|
||||
"2026-06-03T00:00:00+00:00",
|
||||
usage=LLMUsage.reported(input_tokens=10, output_tokens=2),
|
||||
)
|
||||
store.record(first_call)
|
||||
kwargs = {"now": datetime(2026, 6, 3, 12, tzinfo=timezone.utc)}
|
||||
|
||||
first = store.usage_payload(**kwargs)
|
||||
first["days"].clear()
|
||||
cached = store.usage_payload(**kwargs)
|
||||
assert cached["total_tokens"] == 12
|
||||
assert cached["days"]
|
||||
|
||||
store.record(first_call)
|
||||
refreshed = store.usage_payload(**kwargs)
|
||||
assert refreshed["total_tokens"] == 24
|
||||
|
||||
|
||||
def test_usage_payload_cache_is_invalidated_when_connection_pid_changes(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
path = tmp_path / "llm_usage.sqlite3"
|
||||
store = LLMUsageStore(path)
|
||||
other_store = LLMUsageStore(path)
|
||||
first_call = _call(
|
||||
"2026-06-03T00:00:00+00:00",
|
||||
usage=LLMUsage.reported(input_tokens=1, output_tokens=0),
|
||||
)
|
||||
kwargs = {"now": datetime(2026, 6, 3, 12, tzinfo=timezone.utc)}
|
||||
store.record(first_call)
|
||||
assert store.usage_payload(**kwargs)["total_tokens"] == 1
|
||||
|
||||
other_store.record(_call(
|
||||
"2026-06-03T00:01:00+00:00",
|
||||
usage=LLMUsage.reported(input_tokens=2, output_tokens=0),
|
||||
))
|
||||
store._connection_pid = -1
|
||||
|
||||
assert store.usage_payload(**kwargs)["total_tokens"] == 3
|
||||
other_store.close()
|
||||
|
||||
|
||||
def test_usage_query_does_not_hold_writer_lock(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
store = LLMUsageStore(tmp_path / "llm_usage.sqlite3")
|
||||
call = _call(
|
||||
"2026-06-03T00:00:00+00:00",
|
||||
usage=LLMUsage.reported(input_tokens=10, output_tokens=2),
|
||||
)
|
||||
store.record(call)
|
||||
query_started = threading.Event()
|
||||
release_query = threading.Event()
|
||||
record_finished = threading.Event()
|
||||
original_daily_rows = store._daily_rows
|
||||
|
||||
def slow_daily_rows(**kwargs: Any):
|
||||
query_started.set()
|
||||
assert release_query.wait(timeout=2)
|
||||
return original_daily_rows(**kwargs)
|
||||
|
||||
def record_call() -> None:
|
||||
store.record(call)
|
||||
record_finished.set()
|
||||
|
||||
monkeypatch.setattr(store, "_daily_rows", slow_daily_rows)
|
||||
query_thread = threading.Thread(target=lambda: store.usage_payload(
|
||||
now=datetime(2026, 6, 3, 12, tzinfo=timezone.utc),
|
||||
))
|
||||
record_thread = threading.Thread(target=record_call)
|
||||
|
||||
query_thread.start()
|
||||
assert query_started.wait(timeout=2)
|
||||
record_thread.start()
|
||||
try:
|
||||
assert record_finished.wait(timeout=0.5)
|
||||
finally:
|
||||
release_query.set()
|
||||
query_thread.join(timeout=2)
|
||||
record_thread.join(timeout=2)
|
||||
|
||||
assert not query_thread.is_alive()
|
||||
assert not record_thread.is_alive()
|
||||
@@ -0,0 +1,180 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Iterator
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from nanobot.llm_usage.context import llm_usage_source
|
||||
from nanobot.llm_usage.models import LLMCallRecord
|
||||
from nanobot.providers.base import LLMProvider, LLMResponse, LLMUsage
|
||||
from nanobot.providers.fallback_provider import FallbackProvider
|
||||
|
||||
|
||||
class _SequenceProvider(LLMProvider):
|
||||
_CHAT_RETRY_DELAYS = (0,)
|
||||
|
||||
def __init__(self, responses: Iterator[LLMResponse]) -> None:
|
||||
super().__init__(provider_name="test-provider")
|
||||
self._responses = responses
|
||||
|
||||
async def chat(self, **_kwargs: object) -> LLMResponse:
|
||||
return next(self._responses)
|
||||
|
||||
def get_default_model(self) -> str:
|
||||
return "test-model"
|
||||
|
||||
|
||||
class _NoRetryProvider(_SequenceProvider):
|
||||
_CHAT_RETRY_DELAYS = ()
|
||||
|
||||
|
||||
class _BlockingProvider(LLMProvider):
|
||||
async def chat(self, **_kwargs: object) -> LLMResponse:
|
||||
await asyncio.Event().wait()
|
||||
raise AssertionError("unreachable")
|
||||
|
||||
async def chat_stream(self, **_kwargs: object) -> LLMResponse:
|
||||
await asyncio.Event().wait()
|
||||
raise AssertionError("unreachable")
|
||||
|
||||
def get_default_model(self) -> str:
|
||||
return "test-model"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_observer_receives_every_retry_attempt() -> None:
|
||||
provider = _SequenceProvider(
|
||||
iter(
|
||||
[
|
||||
LLMResponse(
|
||||
content="temporary failure",
|
||||
finish_reason="error",
|
||||
error_kind="timeout",
|
||||
),
|
||||
LLMResponse(
|
||||
content="ok",
|
||||
usage=LLMUsage.reported(
|
||||
input_tokens=100,
|
||||
output_tokens=20,
|
||||
cache_read_tokens=60,
|
||||
),
|
||||
),
|
||||
]
|
||||
)
|
||||
)
|
||||
events: list[LLMCallRecord] = []
|
||||
provider.set_llm_call_observer(events.append)
|
||||
|
||||
with llm_usage_source("api"):
|
||||
response = await provider.chat_with_retry(
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
model="selected-model",
|
||||
)
|
||||
|
||||
assert response.finish_reason == "stop"
|
||||
assert len(events) == 2
|
||||
assert [event.finish_reason for event in events] == ["error", "stop"]
|
||||
assert all(event.provider == "test-provider" for event in events)
|
||||
assert all(event.model == "selected-model" for event in events)
|
||||
assert all(event.source == "api" for event in events)
|
||||
assert events[1].usage is not None
|
||||
assert events[1].usage.cache_read_tokens == 60
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_observer_estimates_missing_success_usage_without_storing_content() -> None:
|
||||
provider = _SequenceProvider(iter([LLMResponse(content="hello")]))
|
||||
events: list[LLMCallRecord] = []
|
||||
provider.set_llm_call_observer(events.append)
|
||||
|
||||
response = await provider.chat_with_retry(
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
)
|
||||
|
||||
assert response.usage is not None
|
||||
assert response.usage.source == "estimated"
|
||||
assert events[0].usage == response.usage
|
||||
assert "content" not in LLMCallRecord.__dataclass_fields__
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_observer_failure_never_breaks_provider_call() -> None:
|
||||
provider = _SequenceProvider(iter([LLMResponse(content="ok")]))
|
||||
|
||||
def _fail(_event: LLMCallRecord) -> None:
|
||||
raise RuntimeError("disk unavailable")
|
||||
|
||||
provider.set_llm_call_observer(_fail)
|
||||
response = await provider.chat_with_retry(
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
)
|
||||
|
||||
assert response.content == "ok"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("stream", [False, True])
|
||||
async def test_observer_records_cancelled_provider_attempt(stream: bool) -> None:
|
||||
provider = _BlockingProvider(provider_name="blocking-provider")
|
||||
events: list[LLMCallRecord] = []
|
||||
provider.set_llm_call_observer(events.append)
|
||||
call = provider.chat_stream_with_retry if stream else provider.chat_with_retry
|
||||
|
||||
with pytest.raises(asyncio.TimeoutError):
|
||||
await asyncio.wait_for(
|
||||
call(messages=[{"role": "user", "content": "hello"}]),
|
||||
timeout=0.01,
|
||||
)
|
||||
|
||||
assert len(events) == 1
|
||||
assert events[0].finish_reason == "cancelled"
|
||||
assert events[0].error_kind == "cancelled"
|
||||
assert events[0].usage is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fallback_provider_propagates_observer_to_every_leaf() -> None:
|
||||
primary = _NoRetryProvider(
|
||||
iter(
|
||||
[
|
||||
LLMResponse(
|
||||
content="primary unavailable",
|
||||
finish_reason="error",
|
||||
error_kind="timeout",
|
||||
)
|
||||
]
|
||||
)
|
||||
)
|
||||
fallback = _SequenceProvider(
|
||||
iter(
|
||||
[
|
||||
LLMResponse(
|
||||
content="fallback ok",
|
||||
usage=LLMUsage.reported(input_tokens=12, output_tokens=3),
|
||||
)
|
||||
]
|
||||
)
|
||||
)
|
||||
preset = SimpleNamespace(
|
||||
model="fallback-model",
|
||||
max_tokens=256,
|
||||
temperature=0.2,
|
||||
reasoning_effort=None,
|
||||
context_window_tokens=4_096,
|
||||
)
|
||||
provider = FallbackProvider(primary, [preset], lambda _preset: fallback)
|
||||
events: list[LLMCallRecord] = []
|
||||
provider.set_llm_call_observer(events.append)
|
||||
|
||||
response = await provider.chat_with_retry(
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
model="primary-model",
|
||||
)
|
||||
|
||||
assert response.content == "fallback ok"
|
||||
assert [(event.model, event.finish_reason) for event in events] == [
|
||||
("primary-model", "error"),
|
||||
("fallback-model", "stop"),
|
||||
]
|
||||
@@ -2,6 +2,7 @@ from __future__ import annotations
|
||||
|
||||
import builtins
|
||||
import json
|
||||
import time
|
||||
from types import SimpleNamespace
|
||||
|
||||
import httpx
|
||||
@@ -9,6 +10,8 @@ import pytest
|
||||
|
||||
from nanobot.config.loader import load_config, save_config
|
||||
from nanobot.config.schema import Config, InlineFallbackConfig, ModelPresetConfig
|
||||
from nanobot.llm_usage import get_llm_usage_store
|
||||
from nanobot.llm_usage.models import LLMCallRecord
|
||||
from nanobot.providers.base import LLMUsage
|
||||
from nanobot.providers.registry import find_by_name
|
||||
from nanobot.session.manager import SessionManager
|
||||
@@ -1463,14 +1466,16 @@ def test_settings_payload_includes_token_usage_summary(
|
||||
config = Config()
|
||||
save_config(config, config_path)
|
||||
monkeypatch.setattr("nanobot.config.loader._current_config_path", config_path)
|
||||
monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui")
|
||||
|
||||
from nanobot.webui.token_usage import record_token_usage
|
||||
|
||||
record_token_usage(
|
||||
LLMUsage.reported(input_tokens=10, output_tokens=5),
|
||||
timezone_name=config.agents.defaults.timezone,
|
||||
)
|
||||
get_llm_usage_store().record(LLMCallRecord(
|
||||
started_at_ms=int(time.time() * 1000),
|
||||
duration_ms=1,
|
||||
provider="openai",
|
||||
model="gpt-5",
|
||||
source="user",
|
||||
stream=False,
|
||||
finish_reason="stop",
|
||||
usage=LLMUsage.reported(input_tokens=10, output_tokens=5),
|
||||
))
|
||||
|
||||
payload = settings_payload()
|
||||
|
||||
@@ -1491,14 +1496,16 @@ def test_settings_usage_payload_returns_lightweight_token_usage(
|
||||
config = Config()
|
||||
save_config(config, config_path)
|
||||
monkeypatch.setattr("nanobot.config.loader._current_config_path", config_path)
|
||||
monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui")
|
||||
|
||||
from nanobot.webui.token_usage import record_token_usage
|
||||
|
||||
record_token_usage(
|
||||
LLMUsage.reported(input_tokens=20, output_tokens=2),
|
||||
timezone_name=config.agents.defaults.timezone,
|
||||
)
|
||||
get_llm_usage_store().record(LLMCallRecord(
|
||||
started_at_ms=int(time.time() * 1000),
|
||||
duration_ms=1,
|
||||
provider="openai",
|
||||
model="gpt-5",
|
||||
source="user",
|
||||
stream=False,
|
||||
finish_reason="stop",
|
||||
usage=LLMUsage.reported(input_tokens=20, output_tokens=2),
|
||||
))
|
||||
|
||||
payload = settings_usage_payload()
|
||||
|
||||
|
||||
@@ -2,6 +2,7 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import threading
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
@@ -98,6 +99,45 @@ async def test_mcp_list_serializes_local_runtime_failure_snapshot(tmp_path) -> N
|
||||
assert snapshot_calls == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_usage_query_runs_off_the_event_loop(monkeypatch) -> None:
|
||||
calling_thread = threading.get_ident()
|
||||
worker_threads: list[int] = []
|
||||
|
||||
def usage_payload(**_kwargs):
|
||||
worker_threads.append(threading.get_ident())
|
||||
return {"days": []}
|
||||
|
||||
monkeypatch.setattr("nanobot.webui.settings_routes.settings_usage_payload", usage_payload)
|
||||
request = SimpleNamespace(path="/api/settings/usage", headers=Headers())
|
||||
|
||||
response = await _router().dispatch(None, request, request.path)
|
||||
|
||||
assert response is not None
|
||||
assert response.status_code == 200
|
||||
assert worker_threads and worker_threads[0] != calling_thread
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_full_settings_query_runs_off_the_event_loop(monkeypatch) -> None:
|
||||
calling_thread = threading.get_ident()
|
||||
worker_threads: list[int] = []
|
||||
router = _router()
|
||||
|
||||
def settings_response():
|
||||
worker_threads.append(threading.get_ident())
|
||||
return http_json_response({"ok": True})
|
||||
|
||||
monkeypatch.setattr(router, "_handle_settings", settings_response)
|
||||
request = SimpleNamespace(path="/api/settings", headers=Headers())
|
||||
|
||||
response = await router.dispatch(None, request, request.path)
|
||||
|
||||
assert response is not None
|
||||
assert response.status_code == 200
|
||||
assert worker_threads and worker_threads[0] != calling_thread
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_reload_callback_is_bounded(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
|
||||
@@ -1,319 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from nanobot.agent.hook import AgentHookContext
|
||||
from nanobot.providers.base import LLMUsage
|
||||
from nanobot.webui.token_usage import (
|
||||
TokenUsageHook,
|
||||
read_token_usage_state,
|
||||
record_response_token_usage,
|
||||
record_token_usage,
|
||||
token_usage_payload,
|
||||
write_token_usage_state,
|
||||
)
|
||||
|
||||
|
||||
def _write_state(tmp_path, days: dict) -> None:
|
||||
state_dir = tmp_path / "webui"
|
||||
state_dir.mkdir(parents=True, exist_ok=True)
|
||||
(state_dir / "token-usage.json").write_text(
|
||||
json.dumps({"schema_version": 2, "days": days}), encoding="utf-8"
|
||||
)
|
||||
|
||||
|
||||
def test_payload_tolerates_malformed_persisted_day_keys(tmp_path, monkeypatch) -> None:
|
||||
"""Day keys that are not real dates must not break settings payloads.
|
||||
|
||||
normalize_token_usage_state only length-checks day keys, so a hand-edited
|
||||
10-char key survives reads and atomic rewrites; token_usage_payload then
|
||||
parsed it with an unguarded fromisoformat, failing every /api/settings and
|
||||
/api/settings/usage request until the file was fixed by hand.
|
||||
"""
|
||||
monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui")
|
||||
_write_state(tmp_path, {
|
||||
"not-a-dat3": {"total_tokens": 7, "requests": 1},
|
||||
"2026-13-01": {"total_tokens": 9, "requests": 1},
|
||||
"2026-06-02": {"total_tokens": 5, "requests": 1},
|
||||
})
|
||||
|
||||
payload = token_usage_payload(
|
||||
timezone_name="UTC",
|
||||
now=datetime(2026, 6, 3, 12, 0, tzinfo=timezone.utc),
|
||||
)
|
||||
|
||||
assert payload["total_tokens"] == 5
|
||||
assert payload["total_tokens_30d"] == 5
|
||||
assert payload["requests_30d"] == 1
|
||||
assert payload["active_days_30d"] == 1
|
||||
|
||||
|
||||
def test_record_scrubs_malformed_day_keys(tmp_path, monkeypatch) -> None:
|
||||
"""Rewrites drop malformed day keys instead of persisting them forever."""
|
||||
monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui")
|
||||
_write_state(tmp_path, {
|
||||
"not-a-dat3": {"total_tokens": 7, "requests": 1},
|
||||
"2026-06-02": {"total_tokens": 5, "requests": 1},
|
||||
})
|
||||
|
||||
record_token_usage(
|
||||
LLMUsage.reported(input_tokens=1, output_tokens=1),
|
||||
timezone_name="UTC",
|
||||
now=datetime(2026, 6, 3, 12, 0, tzinfo=timezone.utc),
|
||||
)
|
||||
|
||||
raw = json.loads((tmp_path / "webui" / "token-usage.json").read_text(encoding="utf-8"))
|
||||
assert "not-a-dat3" not in raw["days"]
|
||||
assert "2026-06-02" in raw["days"]
|
||||
assert "2026-06-03" in raw["days"]
|
||||
|
||||
|
||||
def test_record_token_usage_aggregates_by_local_day(tmp_path, monkeypatch) -> None:
|
||||
monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui")
|
||||
|
||||
record_token_usage(
|
||||
LLMUsage.reported(
|
||||
input_tokens=100,
|
||||
output_tokens=40,
|
||||
cache_read_tokens=20,
|
||||
),
|
||||
timezone_name="Asia/Shanghai",
|
||||
now=datetime(2026, 6, 2, 18, 0, tzinfo=timezone.utc),
|
||||
)
|
||||
record_token_usage(
|
||||
LLMUsage.reported(input_tokens=10, output_tokens=5),
|
||||
timezone_name="Asia/Shanghai",
|
||||
now=datetime(2026, 6, 2, 19, 0, tzinfo=timezone.utc),
|
||||
)
|
||||
|
||||
payload = token_usage_payload(
|
||||
timezone_name="Asia/Shanghai",
|
||||
now=datetime(2026, 6, 3, 12, 0, tzinfo=timezone.utc),
|
||||
)
|
||||
|
||||
assert payload["total_tokens_30d"] == 155
|
||||
assert payload["active_days_30d"] == 1
|
||||
assert payload["requests_30d"] == 2
|
||||
assert payload["days"] == [
|
||||
{
|
||||
"date": "2026-06-03",
|
||||
"input_tokens": 110,
|
||||
"output_tokens": 45,
|
||||
"cache_read_tokens": 20,
|
||||
"cache_write_tokens": 0,
|
||||
"cache_read_observed_input_tokens": 100,
|
||||
"cache_write_observed_input_tokens": 0,
|
||||
"total_tokens": 155,
|
||||
"reported_tokens": 155,
|
||||
"estimated_tokens": 0,
|
||||
"requests": 2,
|
||||
"reported_requests": 2,
|
||||
"estimated_requests": 0,
|
||||
"sources": {
|
||||
"user": {
|
||||
"input_tokens": 110,
|
||||
"output_tokens": 45,
|
||||
"cache_read_tokens": 20,
|
||||
"cache_write_tokens": 0,
|
||||
"cache_read_observed_input_tokens": 100,
|
||||
"cache_write_observed_input_tokens": 0,
|
||||
"total_tokens": 155,
|
||||
"reported_tokens": 155,
|
||||
"estimated_tokens": 0,
|
||||
"requests": 2,
|
||||
"reported_requests": 2,
|
||||
"estimated_requests": 0,
|
||||
}
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_cache_observation_denominators_distinguish_missing_from_zero(
|
||||
tmp_path,
|
||||
monkeypatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui")
|
||||
now = datetime(2026, 6, 3, tzinfo=timezone.utc)
|
||||
|
||||
record_token_usage(
|
||||
LLMUsage.reported(input_tokens=100, output_tokens=10),
|
||||
source="user",
|
||||
now=now,
|
||||
)
|
||||
record_token_usage(
|
||||
LLMUsage.reported(
|
||||
input_tokens=40,
|
||||
output_tokens=5,
|
||||
cache_read_tokens=0,
|
||||
cache_write_tokens=0,
|
||||
),
|
||||
source="dream",
|
||||
now=now,
|
||||
)
|
||||
|
||||
row = token_usage_payload(now=now)["days"][0]
|
||||
|
||||
assert row["cache_read_tokens"] == 0
|
||||
assert row["cache_write_tokens"] == 0
|
||||
assert row["cache_read_observed_input_tokens"] == 40
|
||||
assert row["cache_write_observed_input_tokens"] == 40
|
||||
assert row["sources"]["user"]["cache_read_observed_input_tokens"] == 0
|
||||
assert row["sources"]["user"]["cache_write_observed_input_tokens"] == 0
|
||||
assert row["sources"]["dream"]["cache_read_observed_input_tokens"] == 40
|
||||
assert row["sources"]["dream"]["cache_write_observed_input_tokens"] == 40
|
||||
|
||||
|
||||
def _retention_state(sources: tuple[str, ...], *, day_count: int = 400) -> dict:
|
||||
start = datetime(2025, 1, 1, tzinfo=timezone.utc)
|
||||
source_usage = {
|
||||
"input_tokens": 100,
|
||||
"output_tokens": 10,
|
||||
"total_tokens": 110,
|
||||
"reported_tokens": 110,
|
||||
"requests": 1,
|
||||
"reported_requests": 1,
|
||||
}
|
||||
days = {}
|
||||
for offset in range(day_count):
|
||||
day = (start + timedelta(days=offset)).date().isoformat()
|
||||
days[day] = {
|
||||
"input_tokens": 100 * len(sources),
|
||||
"output_tokens": 10 * len(sources),
|
||||
"total_tokens": 110 * len(sources),
|
||||
"reported_tokens": 110 * len(sources),
|
||||
"requests": len(sources),
|
||||
"reported_requests": len(sources),
|
||||
"sources": {source: dict(source_usage) for source in sources},
|
||||
}
|
||||
return {"schema_version": 2, "days": days}
|
||||
|
||||
|
||||
def test_write_compact_state_keeps_400_days_with_two_sources(tmp_path, monkeypatch) -> None:
|
||||
monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui")
|
||||
|
||||
written = write_token_usage_state(_retention_state(("user", "api")))
|
||||
persisted = (tmp_path / "webui" / "token-usage.json").read_bytes()
|
||||
|
||||
assert len(written["days"]) == 400
|
||||
assert len(persisted) <= 512 * 1024
|
||||
assert persisted.endswith(b"\n")
|
||||
assert json.loads(persisted) == written
|
||||
|
||||
|
||||
def test_write_prunes_only_oldest_days_to_fit_byte_budget(tmp_path, monkeypatch) -> None:
|
||||
monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui")
|
||||
sources = ("user", "api", "cron", "dream", "system")
|
||||
raw = _retention_state(sources)
|
||||
all_dates = list(raw["days"])
|
||||
|
||||
written = write_token_usage_state(raw)
|
||||
retained_dates = list(written["days"])
|
||||
persisted = (tmp_path / "webui" / "token-usage.json").read_bytes()
|
||||
|
||||
assert 1 <= len(retained_dates) < len(all_dates)
|
||||
assert retained_dates == all_dates[-len(retained_dates) :]
|
||||
assert retained_dates[-1] == all_dates[-1]
|
||||
assert all(set(row["sources"]) == set(sources) for row in written["days"].values())
|
||||
assert len(persisted) <= 512 * 1024
|
||||
assert read_token_usage_state() == written
|
||||
|
||||
|
||||
def test_write_raises_when_latest_day_alone_exceeds_byte_budget(
|
||||
tmp_path,
|
||||
monkeypatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui")
|
||||
monkeypatch.setattr("nanobot.webui.token_usage._MAX_STATE_FILE_BYTES", 256)
|
||||
|
||||
with pytest.raises(ValueError, match="latest token usage day exceeds"):
|
||||
write_token_usage_state(_retention_state(("user", "api"), day_count=1))
|
||||
|
||||
assert not (tmp_path / "webui" / "token-usage.json").exists()
|
||||
|
||||
|
||||
def test_record_token_usage_skips_empty_usage(tmp_path, monkeypatch) -> None:
|
||||
monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui")
|
||||
|
||||
record_token_usage(LLMUsage.reported(input_tokens=0, output_tokens=0))
|
||||
|
||||
payload = token_usage_payload(now=datetime(2026, 6, 3, tzinfo=timezone.utc))
|
||||
assert payload["days"] == []
|
||||
assert payload["total_tokens_30d"] == 0
|
||||
|
||||
|
||||
def test_record_token_usage_keeps_estimated_split(tmp_path, monkeypatch) -> None:
|
||||
monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui")
|
||||
|
||||
record_token_usage(
|
||||
LLMUsage.estimated(input_tokens=100, output_tokens=25),
|
||||
now=datetime(2026, 6, 3, tzinfo=timezone.utc),
|
||||
)
|
||||
|
||||
payload = token_usage_payload(now=datetime(2026, 6, 3, tzinfo=timezone.utc))
|
||||
|
||||
assert payload["days"][0]["total_tokens"] == 125
|
||||
assert payload["days"][0]["reported_tokens"] == 0
|
||||
assert payload["days"][0]["estimated_tokens"] == 125
|
||||
assert payload["days"][0]["estimated_requests"] == 1
|
||||
|
||||
|
||||
def test_record_token_usage_keeps_source_breakdown(tmp_path, monkeypatch) -> None:
|
||||
monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui")
|
||||
|
||||
record_token_usage(
|
||||
LLMUsage.reported(input_tokens=100, output_tokens=25, total_tokens=175),
|
||||
source="user",
|
||||
now=datetime(2026, 6, 3, tzinfo=timezone.utc),
|
||||
)
|
||||
record_token_usage(
|
||||
LLMUsage.reported(input_tokens=20, output_tokens=5),
|
||||
source="dream",
|
||||
now=datetime(2026, 6, 3, tzinfo=timezone.utc),
|
||||
)
|
||||
|
||||
payload = token_usage_payload(now=datetime(2026, 6, 3, tzinfo=timezone.utc))
|
||||
row = payload["days"][0]
|
||||
|
||||
assert row["total_tokens"] == 200
|
||||
assert row["sources"]["user"]["total_tokens"] == 175
|
||||
assert row["sources"]["user"]["requests"] == 1
|
||||
assert row["sources"]["dream"]["total_tokens"] == 25
|
||||
assert row["sources"]["dream"]["requests"] == 1
|
||||
|
||||
|
||||
def test_record_response_token_usage_uses_response_usage(tmp_path, monkeypatch) -> None:
|
||||
monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui")
|
||||
monkeypatch.setattr("nanobot.webui.token_usage._local_day", lambda *_, **__: "2026-06-03")
|
||||
|
||||
record_response_token_usage(
|
||||
SimpleNamespace(usage=LLMUsage.reported(input_tokens=20, output_tokens=5)),
|
||||
source="dream",
|
||||
)
|
||||
|
||||
payload = token_usage_payload(now=datetime(2026, 6, 3, tzinfo=timezone.utc))
|
||||
assert payload["days"][0]["sources"]["dream"]["total_tokens"] == 25
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_token_usage_hook_classifies_source_from_session_key(tmp_path, monkeypatch) -> None:
|
||||
monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui")
|
||||
monkeypatch.setattr("nanobot.webui.token_usage._local_day", lambda *_, **__: "2026-06-03")
|
||||
|
||||
hook = TokenUsageHook()
|
||||
await hook.after_iteration(
|
||||
AgentHookContext(
|
||||
iteration=0,
|
||||
messages=[],
|
||||
session_key="cron:drink-water",
|
||||
usage=LLMUsage.reported(input_tokens=10, output_tokens=5),
|
||||
)
|
||||
)
|
||||
|
||||
payload = token_usage_payload(now=datetime(2026, 6, 3, tzinfo=timezone.utc))
|
||||
|
||||
assert payload["days"][0]["sources"]["cron"]["total_tokens"] == 15
|
||||
@@ -724,6 +724,13 @@ export interface SettingsPayload {
|
||||
requests: number;
|
||||
reported_requests?: number;
|
||||
estimated_requests?: number;
|
||||
successful_requests?: number;
|
||||
failed_requests?: number;
|
||||
generation_ms?: number;
|
||||
measured_output_tokens?: number;
|
||||
ttft_ms?: number;
|
||||
timed_requests?: number;
|
||||
duration_ms?: number;
|
||||
sources?: Record<
|
||||
"user" | "api" | "cron" | "dream" | "system" | string,
|
||||
{
|
||||
@@ -739,6 +746,13 @@ export interface SettingsPayload {
|
||||
requests: number;
|
||||
reported_requests?: number;
|
||||
estimated_requests?: number;
|
||||
successful_requests?: number;
|
||||
failed_requests?: number;
|
||||
generation_ms?: number;
|
||||
measured_output_tokens?: number;
|
||||
ttft_ms?: number;
|
||||
timed_requests?: number;
|
||||
duration_ms?: number;
|
||||
}
|
||||
>;
|
||||
}>;
|
||||
@@ -750,6 +764,35 @@ export interface SettingsPayload {
|
||||
longest_streak_days: number;
|
||||
active_days_30d: number;
|
||||
requests_30d: number;
|
||||
failed_requests_30d?: number;
|
||||
reported_tokens_30d?: number;
|
||||
estimated_tokens_30d?: number;
|
||||
cache_read_tokens_30d?: number;
|
||||
cache_read_observed_input_tokens_30d?: number;
|
||||
cache_read_rate_30d?: number | null;
|
||||
providers_30d?: Array<{
|
||||
provider: string;
|
||||
model: string;
|
||||
input_tokens: number;
|
||||
output_tokens: number;
|
||||
cache_read_tokens: number;
|
||||
cache_write_tokens: number;
|
||||
cache_read_observed_input_tokens: number;
|
||||
cache_write_observed_input_tokens: number;
|
||||
total_tokens: number;
|
||||
reported_tokens: number;
|
||||
estimated_tokens: number;
|
||||
requests: number;
|
||||
successful_requests: number;
|
||||
failed_requests: number;
|
||||
reported_requests: number;
|
||||
estimated_requests: number;
|
||||
generation_ms: number;
|
||||
measured_output_tokens: number;
|
||||
ttft_ms: number;
|
||||
timed_requests: number;
|
||||
duration_ms: number;
|
||||
}>;
|
||||
updated_at?: string | null;
|
||||
};
|
||||
advanced: {
|
||||
|
||||
Reference in New Issue
Block a user