mirror/nanobot
1
0
Fork 0
mirror of https://github.com/HKUDS/nanobot.git synced 2026-09-04 02:01:48 +03:00
Code Issues Packages Projects Releases Wiki Activity

refactor(providers): define typed usage contract

Browse Source
This commit is contained in:
chengyongru
2026-08-25 01:04:25 +08:00
committed by chengyongru
parent 89c94d8744
commit 9895c23cb5
84 changed files with 1643 additions and 726 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/webui/test_session_context.py
+4 -10
View File
@@ -1,3 +1,4 @@
from nanobot.providers.base import LLMUsage
from nanobot.session import Session
from nanobot.utils.helpers import estimate_message_tokens
from nanobot.webui.session_context import session_context_payload
@@ -58,19 +59,12 @@ def test_session_context_tolerates_untrusted_summary_metadata() -> None:
def test_session_context_sanitizes_usage_metadata() -> None:
usage = LLMUsage.reported(input_tokens=120, output_tokens=8, total_tokens=175)
session = Session(
key="websocket:context",
metadata={
"_last_usage": {
"prompt_tokens": 120,
"completion_tokens": 8,
"negative": -1,
"boolean": True,
"text": "invalid",
}
},
metadata={"_last_usage": usage.to_dict()},
)
payload = session_context_payload(session)
assert payload["last_usage"] == {"prompt_tokens": 120, "completion_tokens": 8}
assert payload["last_usage"] == usage.to_dict()
tests/webui/test_settings_api.py
+3 -2
View File
@@ -9,6 +9,7 @@ import pytest
from nanobot.config.loader import load_config, save_config
from nanobot.config.schema import Config, InlineFallbackConfig, ModelPresetConfig
from nanobot.providers.base import LLMUsage
from nanobot.providers.registry import find_by_name
from nanobot.session.manager import SessionManager
from nanobot.session.model_selection import SESSION_MODEL_PRESET_METADATA_KEY
@@ -1467,7 +1468,7 @@ def test_settings_payload_includes_token_usage_summary(
from nanobot.webui.token_usage import record_token_usage
record_token_usage(
{"prompt_tokens": 10, "completion_tokens": 5},
LLMUsage.reported(input_tokens=10, output_tokens=5),
timezone_name=config.agents.defaults.timezone,
)
@@ -1495,7 +1496,7 @@ def test_settings_usage_payload_returns_lightweight_token_usage(
from nanobot.webui.token_usage import record_token_usage
record_token_usage(
{"prompt_tokens": 20, "completion_tokens": 2},
LLMUsage.reported(input_tokens=20, output_tokens=2),
timezone_name=config.agents.defaults.timezone,
)
tests/webui/test_token_usage.py
+140 -24
View File
@@ -1,17 +1,20 @@
from __future__ import annotations
import json
from datetime import datetime, timezone
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,
)
@@ -19,7 +22,7 @@ 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({"days": days}), encoding="utf-8"
json.dumps({"schema_version": 2, "days": days}), encoding="utf-8"
)
@@ -58,7 +61,7 @@ def test_record_scrubs_malformed_day_keys(tmp_path, monkeypatch) -> None:
})
record_token_usage(
{"prompt_tokens": 1, "completion_tokens": 1},
LLMUsage.reported(input_tokens=1, output_tokens=1),
timezone_name="UTC",
now=datetime(2026, 6, 3, 12, 0, tzinfo=timezone.utc),
)
@@ -73,12 +76,16 @@ def test_record_token_usage_aggregates_by_local_day(tmp_path, monkeypatch) -> No
monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui")
record_token_usage(
{"prompt_tokens": 100, "completion_tokens": 40, "cached_tokens": 20},
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(
{"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
LLMUsage.reported(input_tokens=10, output_tokens=5),
timezone_name="Asia/Shanghai",
now=datetime(2026, 6, 2, 19, 0, tzinfo=timezone.utc),
)
@@ -94,25 +101,31 @@ def test_record_token_usage_aggregates_by_local_day(tmp_path, monkeypatch) -> No
assert payload["days"] == [
{
"date": "2026-06-03",
"prompt_tokens": 110,
"completion_tokens": 45,
"cached_tokens": 20,
"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,
"provider_tokens": 155,
"reported_tokens": 155,
"estimated_tokens": 0,
"requests": 2,
"provider_requests": 2,
"reported_requests": 2,
"estimated_requests": 0,
"sources": {
"user": {
"prompt_tokens": 110,
"completion_tokens": 45,
"cached_tokens": 20,
"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,
"provider_tokens": 155,
"reported_tokens": 155,
"estimated_tokens": 0,
"requests": 2,
"provider_requests": 2,
"reported_requests": 2,
"estimated_requests": 0,
}
},
@@ -120,10 +133,113 @@ def test_record_token_usage_aggregates_by_local_day(tmp_path, monkeypatch) -> No
]
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({"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0})
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"] == []
@@ -134,14 +250,14 @@ 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(
{"prompt_tokens": 100, "completion_tokens": 25, "estimated_tokens": 125},
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]["provider_tokens"] == 0
assert payload["days"][0]["reported_tokens"] == 0
assert payload["days"][0]["estimated_tokens"] == 125
assert payload["days"][0]["estimated_requests"] == 1
@@ -150,12 +266,12 @@ def test_record_token_usage_keeps_source_breakdown(tmp_path, monkeypatch) -> Non
monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui")
record_token_usage(
{"prompt_tokens": 100, "completion_tokens": 25},
LLMUsage.reported(input_tokens=100, output_tokens=25, total_tokens=175),
source="user",
now=datetime(2026, 6, 3, tzinfo=timezone.utc),
)
record_token_usage(
{"prompt_tokens": 20, "completion_tokens": 5},
LLMUsage.reported(input_tokens=20, output_tokens=5),
source="dream",
now=datetime(2026, 6, 3, tzinfo=timezone.utc),
)
@@ -163,8 +279,8 @@ def test_record_token_usage_keeps_source_breakdown(tmp_path, monkeypatch) -> Non
payload = token_usage_payload(now=datetime(2026, 6, 3, tzinfo=timezone.utc))
row = payload["days"][0]
assert row["total_tokens"] == 150
assert row["sources"]["user"]["total_tokens"] == 125
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
@@ -175,7 +291,7 @@ def test_record_response_token_usage_uses_response_usage(tmp_path, monkeypatch)
monkeypatch.setattr("nanobot.webui.token_usage._local_day", lambda *_, **__: "2026-06-03")
record_response_token_usage(
SimpleNamespace(usage={"prompt_tokens": 20, "completion_tokens": 5}),
SimpleNamespace(usage=LLMUsage.reported(input_tokens=20, output_tokens=5)),
source="dream",
)
@@ -194,7 +310,7 @@ async def test_token_usage_hook_classifies_source_from_session_key(tmp_path, mon
iteration=0,
messages=[],
session_key="cron:drink-water",
usage={"prompt_tokens": 10, "completion_tokens": 5},
usage=LLMUsage.reported(input_tokens=10, output_tokens=5),
)
)