mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-28 17:06:05 +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.
189 lines
7.6 KiB
Python
189 lines
7.6 KiB
Python
"""Tests for SubagentTokenCollector callback handler."""
|
|
|
|
from unittest.mock import MagicMock
|
|
from uuid import uuid4
|
|
|
|
from deerflow.subagents.token_collector import SubagentTokenCollector
|
|
|
|
|
|
def _make_llm_response(content="Hello", usage=None, response_metadata=None):
|
|
"""Create a mock LLM response with a message."""
|
|
msg = MagicMock()
|
|
msg.content = content
|
|
msg.usage_metadata = usage
|
|
msg.response_metadata = response_metadata or {}
|
|
|
|
gen = MagicMock()
|
|
gen.message = msg
|
|
|
|
response = MagicMock()
|
|
response.generations = [[gen]]
|
|
return response
|
|
|
|
|
|
def _make_llm_response_from_usages(usages):
|
|
"""Create a mock LLM response with one generation per usage entry."""
|
|
generations = []
|
|
for usage in usages:
|
|
msg = MagicMock()
|
|
msg.content = "chunk"
|
|
msg.usage_metadata = usage
|
|
|
|
gen = MagicMock()
|
|
gen.message = msg
|
|
generations.append([gen])
|
|
|
|
response = MagicMock()
|
|
response.generations = generations
|
|
return response
|
|
|
|
|
|
class TestSubagentTokenCollector:
|
|
def test_collects_usage_from_response(self):
|
|
collector = SubagentTokenCollector(caller="subagent:test")
|
|
usage = {"input_tokens": 100, "output_tokens": 50, "total_tokens": 150}
|
|
collector.on_llm_end(_make_llm_response("Hi", usage=usage), run_id=uuid4())
|
|
records = collector.snapshot_records()
|
|
assert len(records) == 1
|
|
assert records[0]["caller"] == "subagent:test"
|
|
assert records[0]["input_tokens"] == 100
|
|
assert records[0]["output_tokens"] == 50
|
|
assert records[0]["total_tokens"] == 150
|
|
assert "source_run_id" in records[0]
|
|
|
|
def test_collects_model_name_from_response_metadata(self):
|
|
collector = SubagentTokenCollector(caller="subagent:test")
|
|
usage = {"input_tokens": 100, "output_tokens": 50, "total_tokens": 150}
|
|
collector.on_llm_end(
|
|
_make_llm_response("Hi", usage=usage, response_metadata={"model_name": "subagent-model"}),
|
|
run_id=uuid4(),
|
|
)
|
|
|
|
records = collector.snapshot_records()
|
|
|
|
assert len(records) == 1
|
|
assert records[0]["model_name"] == "subagent-model"
|
|
|
|
def test_collects_model_name_from_response_metadata_model_fallback(self):
|
|
collector = SubagentTokenCollector(caller="subagent:test")
|
|
usage = {"input_tokens": 100, "output_tokens": 50, "total_tokens": 150}
|
|
collector.on_llm_end(
|
|
_make_llm_response("Hi", usage=usage, response_metadata={"model": "provider-model"}),
|
|
run_id=uuid4(),
|
|
)
|
|
|
|
records = collector.snapshot_records()
|
|
|
|
assert len(records) == 1
|
|
assert records[0]["model_name"] == "provider-model"
|
|
|
|
def test_total_tokens_zero_uses_input_plus_output(self):
|
|
collector = SubagentTokenCollector(caller="subagent:test")
|
|
usage = {"input_tokens": 200, "output_tokens": 100, "total_tokens": 0}
|
|
collector.on_llm_end(_make_llm_response("Hi", usage=usage), run_id=uuid4())
|
|
records = collector.snapshot_records()
|
|
assert len(records) == 1
|
|
assert records[0]["total_tokens"] == 300
|
|
|
|
def test_total_tokens_missing_uses_input_plus_output(self):
|
|
collector = SubagentTokenCollector(caller="subagent:test")
|
|
usage = {"input_tokens": 30, "output_tokens": 20}
|
|
collector.on_llm_end(_make_llm_response("Hi", usage=usage), run_id=uuid4())
|
|
records = collector.snapshot_records()
|
|
assert len(records) == 1
|
|
assert records[0]["total_tokens"] == 50
|
|
|
|
def test_dedup_same_run_id(self):
|
|
collector = SubagentTokenCollector(caller="subagent:test")
|
|
run_id = uuid4()
|
|
usage = {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}
|
|
collector.on_llm_end(_make_llm_response("A", usage=usage), run_id=run_id)
|
|
collector.on_llm_end(_make_llm_response("A", usage=usage), run_id=run_id)
|
|
records = collector.snapshot_records()
|
|
assert len(records) == 1
|
|
|
|
def test_no_usage_no_record(self):
|
|
collector = SubagentTokenCollector(caller="subagent:test")
|
|
collector.on_llm_end(_make_llm_response("Hi", usage=None), run_id=uuid4())
|
|
records = collector.snapshot_records()
|
|
assert len(records) == 0
|
|
|
|
def test_zero_usage_no_record(self):
|
|
collector = SubagentTokenCollector(caller="subagent:test")
|
|
usage = {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}
|
|
collector.on_llm_end(_make_llm_response("Hi", usage=usage), run_id=uuid4())
|
|
records = collector.snapshot_records()
|
|
assert len(records) == 0
|
|
|
|
def test_skips_empty_generation_and_records_later_usage(self):
|
|
collector = SubagentTokenCollector(caller="subagent:test")
|
|
response = _make_llm_response_from_usages(
|
|
[
|
|
None,
|
|
{"input_tokens": 20, "output_tokens": 10, "total_tokens": 30},
|
|
]
|
|
)
|
|
|
|
collector.on_llm_end(response, run_id=uuid4())
|
|
|
|
records = collector.snapshot_records()
|
|
assert len(records) == 1
|
|
assert records[0]["total_tokens"] == 30
|
|
|
|
def test_snapshot_returns_copy(self):
|
|
collector = SubagentTokenCollector(caller="subagent:test")
|
|
usage = {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}
|
|
collector.on_llm_end(_make_llm_response("Hi", usage=usage), run_id=uuid4())
|
|
snap1 = collector.snapshot_records()
|
|
snap2 = collector.snapshot_records()
|
|
assert snap1 == snap2
|
|
assert snap1 is not snap2
|
|
# Mutating snapshot does not affect internal records
|
|
snap1.append({"source_run_id": "fake"})
|
|
assert len(collector.snapshot_records()) == 1
|
|
|
|
def test_multiple_calls_accumulate(self):
|
|
collector = SubagentTokenCollector(caller="subagent:test")
|
|
usage = {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}
|
|
collector.on_llm_end(_make_llm_response("A", usage=usage), run_id=uuid4())
|
|
collector.on_llm_end(_make_llm_response("B", usage=usage), run_id=uuid4())
|
|
records = collector.snapshot_records()
|
|
assert len(records) == 2
|
|
|
|
def test_different_run_ids_accumulate_separately(self):
|
|
collector = SubagentTokenCollector(caller="subagent:test")
|
|
usage1 = {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}
|
|
usage2 = {"input_tokens": 20, "output_tokens": 10, "total_tokens": 30}
|
|
collector.on_llm_end(_make_llm_response("A", usage=usage1), run_id=uuid4())
|
|
collector.on_llm_end(_make_llm_response("B", usage=usage2), run_id=uuid4())
|
|
records = collector.snapshot_records()
|
|
assert len(records) == 2
|
|
assert records[0]["total_tokens"] == 15
|
|
assert records[1]["total_tokens"] == 30
|
|
|
|
def test_message_without_usage_metadata_skipped(self):
|
|
"""A response where message has no usage_metadata attribute must be skipped."""
|
|
collector = SubagentTokenCollector(caller="subagent:test")
|
|
|
|
msg = MagicMock(spec=[]) # object without usage_metadata
|
|
gen = MagicMock()
|
|
gen.message = msg
|
|
response = MagicMock()
|
|
response.generations = [[gen]]
|
|
|
|
collector.on_llm_end(response, run_id=uuid4())
|
|
records = collector.snapshot_records()
|
|
assert len(records) == 0
|
|
|
|
def test_generation_without_message_skipped(self):
|
|
"""A generation without a message attribute must be skipped."""
|
|
collector = SubagentTokenCollector(caller="subagent:test")
|
|
|
|
gen = MagicMock(spec=[]) # object without message
|
|
response = MagicMock()
|
|
response.generations = [[gen]]
|
|
|
|
collector.on_llm_end(response, run_id=uuid4())
|
|
records = collector.snapshot_records()
|
|
assert len(records) == 0
|