mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-29 09:26:00 +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.
73 lines
2.7 KiB
Python
73 lines
2.7 KiB
Python
"""Callback handler that collects LLM token usage within a subagent.
|
|
|
|
Each subagent execution creates its own collector. After the subagent
|
|
finishes, the collected records are transferred to the parent RunJournal
|
|
via :meth:`RunJournal.record_external_llm_usage_records`.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Mapping
|
|
from typing import Any
|
|
|
|
from langchain_core.callbacks import BaseCallbackHandler
|
|
|
|
|
|
class SubagentTokenCollector(BaseCallbackHandler):
|
|
"""Lightweight callback handler that collects LLM token usage within a subagent."""
|
|
|
|
def __init__(self, caller: str):
|
|
super().__init__()
|
|
self.caller = caller
|
|
self._records: list[dict[str, int | str | None]] = []
|
|
self._counted_run_ids: set[str] = set()
|
|
|
|
def on_llm_end(
|
|
self,
|
|
response: Any,
|
|
*,
|
|
run_id: Any,
|
|
tags: list[str] | None = None,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
rid = str(run_id)
|
|
if rid in self._counted_run_ids:
|
|
return
|
|
|
|
for generation in response.generations:
|
|
for gen in generation:
|
|
if not hasattr(gen, "message"):
|
|
continue
|
|
usage = getattr(gen.message, "usage_metadata", None)
|
|
usage_dict = dict(usage) if usage else {}
|
|
input_tk = usage_dict.get("input_tokens", 0) or 0
|
|
output_tk = usage_dict.get("output_tokens", 0) or 0
|
|
total_tk = usage_dict.get("total_tokens", 0) or 0
|
|
if total_tk <= 0:
|
|
total_tk = input_tk + output_tk
|
|
if total_tk <= 0:
|
|
continue
|
|
# Capture the model that actually produced this response so the
|
|
# parent journal can bucket tokens by real model rather than the
|
|
# lead agent's resolved model
|
|
response_metadata = getattr(gen.message, "response_metadata", None) or {}
|
|
model_name: str | None = None
|
|
if isinstance(response_metadata, Mapping):
|
|
model_name = response_metadata.get("model_name") or response_metadata.get("model")
|
|
self._counted_run_ids.add(rid)
|
|
self._records.append(
|
|
{
|
|
"source_run_id": rid,
|
|
"caller": self.caller,
|
|
"model_name": model_name,
|
|
"input_tokens": input_tk,
|
|
"output_tokens": output_tk,
|
|
"total_tokens": total_tk,
|
|
}
|
|
)
|
|
return
|
|
|
|
def snapshot_records(self) -> list[dict[str, int | str | None]]:
|
|
"""Return a copy of the accumulated usage records."""
|
|
return list(self._records)
|