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

refactor(providers): define typed usage contract

Browse Source
This commit is contained in:
chengyongru
2026-08-25 01:04:25 +08:00
committed by chengyongru
parent 89c94d8744
commit 9895c23cb5
84 changed files with 1643 additions and 726 deletions
Download Patch File Download Diff File Expand all files Collapse all files
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),
)
)