mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-10 00:18:36 +00:00
4669d3c089
* feat(gateway): cache-aware cost accounting + /api/console observability endpoints - Capture prompt-cache hits (usage_metadata.input_token_details.cache_read) in RunJournal and SubagentTokenCollector as a sparse cache_read_tokens key in token_usage_by_model (JSON field — no schema migration; legacy bucket shapes unchanged) - New read-only /api/console router: GET /stats (headline counters), GET /runs (cross-thread paginated history joined with thread titles), GET /usage (zero-filled daily token series + per-model breakdown); user-scoped, 503 on the memory database backend - Optional models[*].pricing (currency, input_per_million, output_per_million, input_cache_hit_per_million) powers real spend estimation; cache-hit input tokens are billed at the hit price (omitted hit price falls back to the miss price as a conservative upper bound); unpriced models yield cost: null - create_chat_model strips the presentation-only pricing block so it never reaches the provider client (unknown kwargs are forwarded into the completion payload and break live calls) - Tests: console router SQLite round-trips, journal/collector cache capture incl. a DeepSeek raw-usage pin test, factory strip regression Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * refactor: address review feedback on cost sum and sparse cache_read_tokens - console.py: replace the walrus-in-generator total-cost sum with an explicit loop (review noted the multi-line form reads ambiguously) - token_collector.py: omit cache_read_tokens from usage records when the provider reported no cache hits, matching the journal's sparse per-model bucket shape; absent is treated as 0 downstream - add a regression test pinning the sparse record shape Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: coffeeFish <codeingforcoffee@users.noreply.github.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
567 lines
22 KiB
Python
567 lines
22 KiB
Python
"""Per-model token usage regression tests (issue #3645).
|
||
|
||
Covers the full path that powers ``GET /api/threads/{id}/token-usage``'s
|
||
``by_model`` field:
|
||
|
||
* ``RunJournal`` capturing each LLM call's real ``response_metadata.model_name``
|
||
for both the lead agent / middleware path (``on_llm_end``) and the subagent
|
||
external-records path (``record_external_llm_usage_records``).
|
||
* ``RunJournal.get_completion_data`` exposing the per-model breakdown so it can
|
||
be threaded into the run store on completion.
|
||
* ``MemoryRunStore`` and ``RunRepository`` (SQLAlchemy) returning the same
|
||
``by_model`` shape from ``aggregate_tokens_by_thread``, with the invariant
|
||
``sum(by_model[*].tokens) == total_tokens``.
|
||
* Legacy rows written before this fix (``token_usage_by_model`` empty) falling
|
||
back to the old ``model_name + total_tokens`` attribution instead of being
|
||
silently dropped.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from unittest.mock import MagicMock
|
||
from uuid import uuid4
|
||
|
||
import pytest
|
||
|
||
from deerflow.persistence.run import RunRepository
|
||
from deerflow.runtime.events.store.memory import MemoryRunEventStore
|
||
from deerflow.runtime.journal import RunJournal
|
||
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Test doubles
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _make_llm_response(*, usage: dict | None, model_name: str | None = "lead-model"):
|
||
"""Build a minimal LLM response carrying the bits journal/collector read."""
|
||
msg = MagicMock()
|
||
msg.type = "ai"
|
||
msg.content = ""
|
||
msg.id = f"msg-{id(msg)}"
|
||
msg.tool_calls = []
|
||
msg.invalid_tool_calls = []
|
||
msg.response_metadata = {} if model_name is None else {"model_name": model_name}
|
||
msg.usage_metadata = usage
|
||
msg.additional_kwargs = {}
|
||
msg.name = None
|
||
msg.model_dump.return_value = {
|
||
"content": "",
|
||
"additional_kwargs": {},
|
||
"response_metadata": msg.response_metadata,
|
||
"type": "ai",
|
||
"name": None,
|
||
"id": msg.id,
|
||
"tool_calls": [],
|
||
"invalid_tool_calls": [],
|
||
"usage_metadata": usage,
|
||
}
|
||
gen = MagicMock()
|
||
gen.message = msg
|
||
response = MagicMock()
|
||
response.generations = [[gen]]
|
||
return response
|
||
|
||
|
||
def _journal() -> RunJournal:
|
||
return RunJournal("r1", "t1", MemoryRunEventStore(), flush_threshold=100)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# RunJournal: per-call model accounting
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestJournalByModel:
|
||
def test_lead_agent_call_lands_on_real_model(self) -> None:
|
||
j = _journal()
|
||
j.on_llm_end(
|
||
_make_llm_response(usage={"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, model_name="lead-model"),
|
||
run_id=uuid4(),
|
||
parent_run_id=None,
|
||
tags=["lead_agent"],
|
||
)
|
||
data = j.get_completion_data()
|
||
assert data["token_usage_by_model"] == {
|
||
"lead-model": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
|
||
}
|
||
assert data["lead_agent_tokens"] == 15
|
||
assert data["total_tokens"] == 15
|
||
|
||
def test_middleware_call_lands_on_its_own_model(self) -> None:
|
||
"""A middleware (e.g. title/summarization) on a different model gets its own bucket."""
|
||
j = _journal()
|
||
j.on_llm_end(
|
||
_make_llm_response(usage={"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, model_name="lead-model"),
|
||
run_id=uuid4(),
|
||
parent_run_id=None,
|
||
tags=["lead_agent"],
|
||
)
|
||
j.on_llm_end(
|
||
_make_llm_response(usage={"input_tokens": 4, "output_tokens": 1, "total_tokens": 5}, model_name="title-model"),
|
||
run_id=uuid4(),
|
||
parent_run_id=None,
|
||
tags=["middleware:title"],
|
||
)
|
||
data = j.get_completion_data()
|
||
assert data["token_usage_by_model"] == {
|
||
"lead-model": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
|
||
"title-model": {"input_tokens": 4, "output_tokens": 1, "total_tokens": 5},
|
||
}
|
||
assert data["lead_agent_tokens"] == 15
|
||
assert data["middleware_tokens"] == 5
|
||
|
||
def test_missing_model_name_falls_back_to_unknown(self) -> None:
|
||
j = _journal()
|
||
j.on_llm_end(
|
||
_make_llm_response(usage={"input_tokens": 3, "output_tokens": 2, "total_tokens": 5}, model_name=None),
|
||
run_id=uuid4(),
|
||
parent_run_id=None,
|
||
tags=["lead_agent"],
|
||
)
|
||
data = j.get_completion_data()
|
||
assert data["token_usage_by_model"] == {
|
||
"unknown": {"input_tokens": 3, "output_tokens": 2, "total_tokens": 5},
|
||
}
|
||
|
||
def test_same_model_aggregates_across_calls(self) -> None:
|
||
j = _journal()
|
||
for _ in range(2):
|
||
j.on_llm_end(
|
||
_make_llm_response(usage={"input_tokens": 7, "output_tokens": 3, "total_tokens": 10}, model_name="lead-model"),
|
||
run_id=uuid4(),
|
||
parent_run_id=None,
|
||
tags=["lead_agent"],
|
||
)
|
||
data = j.get_completion_data()
|
||
assert data["token_usage_by_model"] == {
|
||
"lead-model": {"input_tokens": 14, "output_tokens": 6, "total_tokens": 20},
|
||
}
|
||
|
||
def test_subagent_external_records_attribute_to_real_model(self) -> None:
|
||
"""The fix's headline behavior: subagent on a different model no longer
|
||
steals tokens from the lead model bucket."""
|
||
j = _journal()
|
||
# Lead emits 10 tokens on lead-model.
|
||
j.on_llm_end(
|
||
_make_llm_response(usage={"input_tokens": 6, "output_tokens": 4, "total_tokens": 10}, model_name="lead-model"),
|
||
run_id=uuid4(),
|
||
parent_run_id=None,
|
||
tags=["lead_agent"],
|
||
)
|
||
# Subagent ran on subagent-model and reports 25 tokens via the
|
||
# external-records bridge (the path SubagentTokenCollector uses).
|
||
j.record_external_llm_usage_records(
|
||
[
|
||
{
|
||
"source_run_id": "sub-1",
|
||
"caller": "subagent:general-purpose",
|
||
"model_name": "subagent-model",
|
||
"input_tokens": 15,
|
||
"output_tokens": 10,
|
||
"total_tokens": 25,
|
||
},
|
||
],
|
||
)
|
||
data = j.get_completion_data()
|
||
assert data["token_usage_by_model"] == {
|
||
"lead-model": {"input_tokens": 6, "output_tokens": 4, "total_tokens": 10},
|
||
"subagent-model": {"input_tokens": 15, "output_tokens": 10, "total_tokens": 25},
|
||
}
|
||
assert data["total_tokens"] == 35
|
||
# by_caller stays accurate too.
|
||
assert data["lead_agent_tokens"] == 10
|
||
assert data["subagent_tokens"] == 25
|
||
# Invariant the issue calls out: by_model sums to total_tokens.
|
||
assert sum(b["total_tokens"] for b in data["token_usage_by_model"].values()) == data["total_tokens"]
|
||
|
||
def test_subagent_record_without_model_falls_back_to_unknown(self) -> None:
|
||
j = _journal()
|
||
j.record_external_llm_usage_records(
|
||
[
|
||
{
|
||
"source_run_id": "sub-1",
|
||
"caller": "subagent:bash",
|
||
"input_tokens": 5,
|
||
"output_tokens": 2,
|
||
"total_tokens": 7,
|
||
},
|
||
],
|
||
)
|
||
data = j.get_completion_data()
|
||
assert data["token_usage_by_model"] == {
|
||
"unknown": {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7},
|
||
}
|
||
|
||
def test_on_llm_end_dedup_does_not_double_count_model(self) -> None:
|
||
j = _journal()
|
||
rid = uuid4()
|
||
usage = {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}
|
||
j.on_llm_end(_make_llm_response(usage=usage, model_name="lead-model"), run_id=rid, parent_run_id=None, tags=["lead_agent"])
|
||
# Same langchain run_id firing twice (real callbacks do this) must
|
||
# not inflate either total_tokens or the per-model bucket.
|
||
j.on_llm_end(_make_llm_response(usage=usage, model_name="lead-model"), run_id=rid, parent_run_id=None, tags=["lead_agent"])
|
||
data = j.get_completion_data()
|
||
assert data["total_tokens"] == 15
|
||
assert data["token_usage_by_model"] == {
|
||
"lead-model": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
|
||
}
|
||
|
||
def test_external_records_dedup_does_not_double_count_model(self) -> None:
|
||
j = _journal()
|
||
record = {
|
||
"source_run_id": "sub-1",
|
||
"caller": "subagent:general-purpose",
|
||
"model_name": "subagent-model",
|
||
"input_tokens": 15,
|
||
"output_tokens": 10,
|
||
"total_tokens": 25,
|
||
}
|
||
j.record_external_llm_usage_records([record])
|
||
j.record_external_llm_usage_records([record])
|
||
data = j.get_completion_data()
|
||
assert data["subagent_tokens"] == 25
|
||
assert data["token_usage_by_model"] == {
|
||
"subagent-model": {"input_tokens": 15, "output_tokens": 10, "total_tokens": 25},
|
||
}
|
||
|
||
def test_track_tokens_disabled_keeps_by_model_empty(self) -> None:
|
||
store = MemoryRunEventStore()
|
||
j = RunJournal("r1", "t1", store, track_token_usage=False, flush_threshold=100)
|
||
j.on_llm_end(
|
||
_make_llm_response(usage={"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, model_name="lead-model"),
|
||
run_id=uuid4(),
|
||
parent_run_id=None,
|
||
tags=["lead_agent"],
|
||
)
|
||
j.record_external_llm_usage_records(
|
||
[{"source_run_id": "sub", "caller": "subagent:x", "model_name": "sub-model", "input_tokens": 1, "output_tokens": 1, "total_tokens": 2}],
|
||
)
|
||
data = j.get_completion_data()
|
||
assert data["token_usage_by_model"] == {}
|
||
assert data["total_tokens"] == 0
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Store aggregation: invariants and parity across MemoryRunStore + RunRepository
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
_THREAD = "thread-by-model"
|
||
|
||
|
||
def _completed_run(
|
||
run_id: str,
|
||
*,
|
||
model_name: str | None,
|
||
total_tokens: int,
|
||
lead: int = 0,
|
||
sub: int = 0,
|
||
mw: int = 0,
|
||
by_model: dict | None = None,
|
||
) -> dict:
|
||
"""Shape that both stores accept for completion writes (kwargs to update_run_completion)."""
|
||
return {
|
||
"run_id": run_id,
|
||
"model_name": model_name,
|
||
"completion": {
|
||
"status": "success",
|
||
"total_input_tokens": 0,
|
||
"total_output_tokens": 0,
|
||
"total_tokens": total_tokens,
|
||
"llm_call_count": 1,
|
||
"lead_agent_tokens": lead,
|
||
"subagent_tokens": sub,
|
||
"middleware_tokens": mw,
|
||
"token_usage_by_model": by_model or {},
|
||
"message_count": 0,
|
||
},
|
||
}
|
||
|
||
|
||
async def _seed_run(store, *, run_id: str, model_name: str | None, completion: dict) -> None:
|
||
await store.put(run_id, thread_id=_THREAD, status="pending", model_name=model_name)
|
||
await store.update_run_completion(run_id, **completion)
|
||
|
||
|
||
_RUN_FIXTURES = [
|
||
# 1. Run where subagent and middleware ran on different models than lead.
|
||
_completed_run(
|
||
"run-1",
|
||
model_name="lead-model",
|
||
total_tokens=300,
|
||
lead=100,
|
||
sub=150,
|
||
mw=50,
|
||
by_model={
|
||
"lead-model": {"input_tokens": 60, "output_tokens": 40, "total_tokens": 100},
|
||
"subagent-model": {"input_tokens": 90, "output_tokens": 60, "total_tokens": 150},
|
||
"middleware-model": {"input_tokens": 30, "output_tokens": 20, "total_tokens": 50},
|
||
},
|
||
),
|
||
# 2. Another run, lead on a *different* lead model — exercises multi-run merge.
|
||
_completed_run(
|
||
"run-2",
|
||
model_name="lead-model-b",
|
||
total_tokens=80,
|
||
lead=80,
|
||
by_model={
|
||
"lead-model-b": {"input_tokens": 50, "output_tokens": 30, "total_tokens": 80},
|
||
},
|
||
),
|
||
# 3. Legacy row written before this fix: empty token_usage_by_model. Must
|
||
# fall back to (model_name, total_tokens) instead of disappearing from
|
||
# by_model entirely.
|
||
_completed_run(
|
||
"run-3",
|
||
model_name="legacy-model",
|
||
total_tokens=42,
|
||
lead=42,
|
||
by_model={},
|
||
),
|
||
]
|
||
|
||
|
||
async def _seed_all(store) -> None:
|
||
for fix in _RUN_FIXTURES:
|
||
await _seed_run(store, run_id=fix["run_id"], model_name=fix["model_name"], completion=fix["completion"])
|
||
|
||
|
||
def _assert_aggregate_shape(agg: dict) -> None:
|
||
"""Pin the contract that powers /api/threads/{id}/token-usage."""
|
||
# The headline totals stay the simple SUMs.
|
||
assert agg["total_tokens"] == 300 + 80 + 42
|
||
assert agg["total_runs"] == 3
|
||
assert agg["by_caller"] == {
|
||
"lead_agent": 100 + 80 + 42,
|
||
"subagent": 150,
|
||
"middleware": 50,
|
||
}
|
||
# The core fix: subagent / middleware models show up in by_model with their
|
||
# real tokens; the lead-model bucket is NOT inflated with subagent tokens.
|
||
assert agg["by_model"]["lead-model"] == {"tokens": 100, "runs": 1}
|
||
assert agg["by_model"]["subagent-model"] == {"tokens": 150, "runs": 1}
|
||
assert agg["by_model"]["middleware-model"] == {"tokens": 50, "runs": 1}
|
||
assert agg["by_model"]["lead-model-b"] == {"tokens": 80, "runs": 1}
|
||
# Legacy fallback path — empty token_usage_by_model maps to the row's
|
||
# ``model_name`` with the full total_tokens.
|
||
assert agg["by_model"]["legacy-model"] == {"tokens": 42, "runs": 1}
|
||
# Invariant from issue #3645.
|
||
assert sum(b["tokens"] for b in agg["by_model"].values()) == agg["total_tokens"]
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_memory_store_by_model_invariant_and_fallback():
|
||
store = MemoryRunStore()
|
||
await _seed_all(store)
|
||
agg = await store.aggregate_tokens_by_thread(_THREAD)
|
||
_assert_aggregate_shape(agg)
|
||
|
||
|
||
async def _make_sql_repo(tmp_path):
|
||
from deerflow.persistence.engine import get_session_factory, init_engine
|
||
|
||
url = f"sqlite+aiosqlite:///{tmp_path / 'by-model.db'}"
|
||
await init_engine("sqlite", url=url, sqlite_dir=str(tmp_path))
|
||
return RunRepository(get_session_factory())
|
||
|
||
|
||
async def _close_sql_engine() -> None:
|
||
from deerflow.persistence.engine import close_engine
|
||
|
||
await close_engine()
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_sql_store_by_model_invariant_and_fallback(tmp_path):
|
||
repo = await _make_sql_repo(tmp_path)
|
||
try:
|
||
await _seed_all(repo)
|
||
agg = await repo.aggregate_tokens_by_thread(_THREAD)
|
||
_assert_aggregate_shape(agg)
|
||
finally:
|
||
await _close_sql_engine()
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_memory_and_sql_stores_agree(tmp_path):
|
||
"""Memory and SQL stores must return byte-identical aggregations so
|
||
behavior does not silently diverge based on database.backend choice."""
|
||
mem = MemoryRunStore()
|
||
sql = await _make_sql_repo(tmp_path)
|
||
try:
|
||
await _seed_all(mem)
|
||
await _seed_all(sql)
|
||
mem_agg = await mem.aggregate_tokens_by_thread(_THREAD)
|
||
sql_agg = await sql.aggregate_tokens_by_thread(_THREAD)
|
||
assert mem_agg == sql_agg
|
||
finally:
|
||
await _close_sql_engine()
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_include_active_picks_up_running_progress_snapshot(tmp_path):
|
||
"""``update_run_progress`` must persist ``token_usage_by_model`` so the
|
||
``include_active=true`` view of /token-usage reflects in-flight tokens."""
|
||
repo = await _make_sql_repo(tmp_path)
|
||
try:
|
||
await repo.put("run-active", thread_id=_THREAD, status="pending")
|
||
# Transition to running so update_run_progress' status guard fires.
|
||
await repo.update_status("run-active", "running")
|
||
await repo.update_run_progress(
|
||
"run-active",
|
||
total_tokens=70,
|
||
total_input_tokens=40,
|
||
total_output_tokens=30,
|
||
lead_agent_tokens=70,
|
||
token_usage_by_model={
|
||
"lead-model": {"input_tokens": 40, "output_tokens": 30, "total_tokens": 70},
|
||
},
|
||
)
|
||
# Default (completed-only) excludes running runs.
|
||
completed_only = await repo.aggregate_tokens_by_thread(_THREAD)
|
||
assert completed_only["total_runs"] == 0
|
||
assert completed_only["by_model"] == {}
|
||
|
||
active = await repo.aggregate_tokens_by_thread(_THREAD, include_active=True)
|
||
assert active["total_runs"] == 1
|
||
assert active["by_model"] == {"lead-model": {"tokens": 70, "runs": 1}}
|
||
assert active["total_tokens"] == 70
|
||
finally:
|
||
await _close_sql_engine()
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Prompt-cache-hit accounting (powers cache-aware cost estimation in
|
||
# /api/console): cache_read_tokens is a *sparse* bucket key — present only
|
||
# when a provider reported cache hits — so pre-existing bucket shapes and
|
||
# exact-equality assertions above stay valid.
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestJournalCacheRead:
|
||
def test_cache_read_accumulates_as_sparse_key(self) -> None:
|
||
j = _journal()
|
||
j.on_llm_end(
|
||
_make_llm_response(
|
||
usage={"input_tokens": 100, "output_tokens": 10, "total_tokens": 110, "input_token_details": {"cache_read": 80}},
|
||
model_name="m",
|
||
),
|
||
run_id=uuid4(),
|
||
parent_run_id=None,
|
||
tags=["lead_agent"],
|
||
)
|
||
# Second call without cache hits still accumulates into the same bucket.
|
||
j.on_llm_end(
|
||
_make_llm_response(usage={"input_tokens": 50, "output_tokens": 5, "total_tokens": 55}, model_name="m"),
|
||
run_id=uuid4(),
|
||
parent_run_id=None,
|
||
tags=["lead_agent"],
|
||
)
|
||
data = j.get_completion_data()
|
||
assert data["token_usage_by_model"]["m"] == {
|
||
"input_tokens": 150,
|
||
"output_tokens": 15,
|
||
"total_tokens": 165,
|
||
"cache_read_tokens": 80,
|
||
}
|
||
|
||
def test_bucket_without_cache_hits_keeps_legacy_shape(self) -> None:
|
||
j = _journal()
|
||
j.on_llm_end(
|
||
_make_llm_response(usage={"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, model_name="m"),
|
||
run_id=uuid4(),
|
||
parent_run_id=None,
|
||
tags=["lead_agent"],
|
||
)
|
||
assert j.get_completion_data()["token_usage_by_model"]["m"] == {
|
||
"input_tokens": 10,
|
||
"output_tokens": 5,
|
||
"total_tokens": 15,
|
||
}
|
||
|
||
def test_external_records_carry_cache_read(self) -> None:
|
||
j = _journal()
|
||
j.record_external_llm_usage_records(
|
||
[
|
||
{
|
||
"source_run_id": "sub-1",
|
||
"caller": "subagent:general-purpose",
|
||
"model_name": "sub-m",
|
||
"input_tokens": 40,
|
||
"output_tokens": 10,
|
||
"total_tokens": 50,
|
||
"cache_read_tokens": 25,
|
||
},
|
||
],
|
||
)
|
||
assert j.get_completion_data()["token_usage_by_model"]["sub-m"]["cache_read_tokens"] == 25
|
||
|
||
def test_deepseek_raw_usage_normalizes_to_cache_read(self) -> None:
|
||
"""Pin the DeepSeek chat-completions shape end-to-end: the raw
|
||
``prompt_tokens_details.cached_tokens`` field is what langchain-openai's
|
||
``_create_usage_metadata`` normalizes into
|
||
``input_token_details.cache_read`` (DeepSeek's top-level
|
||
``prompt_cache_hit/miss_tokens`` are redundant aliases LangChain does
|
||
not read), and the journal captures it. The derived cache-miss count
|
||
(input − cache_read) must equal DeepSeek's own
|
||
``prompt_cache_miss_tokens``, which is what cache-aware pricing bills
|
||
at the full input price."""
|
||
from langchain_openai.chat_models.base import _create_usage_metadata
|
||
|
||
raw = {
|
||
"prompt_tokens": 106,
|
||
"completion_tokens": 112,
|
||
"total_tokens": 218,
|
||
"prompt_tokens_details": {"cached_tokens": 64},
|
||
"prompt_cache_hit_tokens": 64,
|
||
"prompt_cache_miss_tokens": 42,
|
||
}
|
||
usage = _create_usage_metadata(raw)
|
||
j = _journal()
|
||
j.on_llm_end(
|
||
_make_llm_response(usage=dict(usage), model_name="deepseek-chat"),
|
||
run_id=uuid4(),
|
||
parent_run_id=None,
|
||
tags=["lead_agent"],
|
||
)
|
||
bucket = j.get_completion_data()["token_usage_by_model"]["deepseek-chat"]
|
||
assert bucket == {
|
||
"input_tokens": 106,
|
||
"output_tokens": 112,
|
||
"total_tokens": 218,
|
||
"cache_read_tokens": 64,
|
||
}
|
||
assert bucket["input_tokens"] - bucket["cache_read_tokens"] == raw["prompt_cache_miss_tokens"]
|
||
|
||
def test_collector_extracts_cache_read_from_usage_metadata(self) -> None:
|
||
from deerflow.subagents.token_collector import SubagentTokenCollector
|
||
|
||
collector = SubagentTokenCollector("subagent:general-purpose")
|
||
collector.on_llm_end(
|
||
_make_llm_response(
|
||
usage={"input_tokens": 30, "output_tokens": 6, "total_tokens": 36, "input_token_details": {"cache_read": 20}},
|
||
model_name="sub-m",
|
||
),
|
||
run_id=uuid4(),
|
||
)
|
||
records = collector.snapshot_records()
|
||
assert len(records) == 1
|
||
assert records[0]["cache_read_tokens"] == 20
|
||
|
||
def test_collector_omits_cache_read_key_when_no_cache_hits(self) -> None:
|
||
from deerflow.subagents.token_collector import SubagentTokenCollector
|
||
|
||
collector = SubagentTokenCollector("subagent:general-purpose")
|
||
collector.on_llm_end(
|
||
_make_llm_response(
|
||
usage={"input_tokens": 30, "output_tokens": 6, "total_tokens": 36},
|
||
model_name="sub-m",
|
||
),
|
||
run_id=uuid4(),
|
||
)
|
||
records = collector.snapshot_records()
|
||
assert len(records) == 1
|
||
# Sparse record shape: no explicit 0 when the provider reported no
|
||
# cache hits (record_external_llm_usage_records treats absent as 0).
|
||
assert "cache_read_tokens" not in records[0]
|