mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-21 21:48:04 +00:00
e7a03e5243
* fix(gateway): attribute token usage to actual models
Capture per-call model names from LLM response metadata for lead, middleware, and subagent calls.
Persist a per-run token_usage_by_model breakdown and aggregate by that map in both SQL and memory stores, with legacy fallback to the run-level model_name for older rows.
Add regression coverage for by_model totals, caller consistency, active progress snapshots, store parity, and SubagentTokenCollector model propagation.
* fix(gateway): harden by-model token aggregation
Use usage.get("total_tokens", 0) when reducing per-model token usage maps so aggregation tolerates partially written or manually edited JSON blobs without changing behavior for journal-written rows.
* docs(gateway): clarify by-model run count semantics
Document that by_model[*].runs counts the number of runs in which a model appeared, so multi-model runs can increment multiple model buckets.
432 lines
16 KiB
Python
432 lines
16 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()
|