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.
145 lines
4.2 KiB
Python
145 lines
4.2 KiB
Python
"""Abstract interface for run metadata storage.
|
|
|
|
RunManager depends on this interface. Implementations:
|
|
- MemoryRunStore: in-memory dict (development, tests)
|
|
- Future: RunRepository backed by SQLAlchemy ORM
|
|
|
|
All methods accept an optional user_id for user isolation.
|
|
When user_id is None, no user filtering is applied (single-user mode).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import abc
|
|
from typing import Any
|
|
|
|
|
|
class RunStore(abc.ABC):
|
|
@abc.abstractmethod
|
|
async def put(
|
|
self,
|
|
run_id: str,
|
|
*,
|
|
thread_id: str,
|
|
assistant_id: str | None = None,
|
|
user_id: str | None = None,
|
|
model_name: str | None = None,
|
|
status: str = "pending",
|
|
multitask_strategy: str = "reject",
|
|
metadata: dict[str, Any] | None = None,
|
|
kwargs: dict[str, Any] | None = None,
|
|
error: str | None = None,
|
|
created_at: str | None = None,
|
|
) -> None:
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
async def get(
|
|
self,
|
|
run_id: str,
|
|
*,
|
|
user_id: str | None = None,
|
|
) -> dict[str, Any] | None:
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
async def list_by_thread(
|
|
self,
|
|
thread_id: str,
|
|
*,
|
|
user_id: str | None = None,
|
|
limit: int = 100,
|
|
) -> list[dict[str, Any]]:
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
async def update_status(
|
|
self,
|
|
run_id: str,
|
|
status: str,
|
|
*,
|
|
error: str | None = None,
|
|
) -> bool | None:
|
|
"""Update a run status.
|
|
|
|
Returns ``False`` when the store can prove no row was updated. Older or
|
|
lightweight stores may return ``None`` when they cannot report rowcount.
|
|
"""
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
async def delete(self, run_id: str) -> None:
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
async def update_model_name(
|
|
self,
|
|
run_id: str,
|
|
model_name: str | None,
|
|
) -> None:
|
|
"""Update the model_name field for an existing run."""
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
async def update_run_completion(
|
|
self,
|
|
run_id: str,
|
|
*,
|
|
status: str,
|
|
total_input_tokens: int = 0,
|
|
total_output_tokens: int = 0,
|
|
total_tokens: int = 0,
|
|
llm_call_count: int = 0,
|
|
lead_agent_tokens: int = 0,
|
|
subagent_tokens: int = 0,
|
|
middleware_tokens: int = 0,
|
|
token_usage_by_model: dict[str, dict[str, int]] | None = None,
|
|
message_count: int = 0,
|
|
last_ai_message: str | None = None,
|
|
first_human_message: str | None = None,
|
|
error: str | None = None,
|
|
) -> bool | None:
|
|
"""Persist final completion fields.
|
|
|
|
Returns ``False`` when the store can prove no row was updated.
|
|
"""
|
|
pass
|
|
|
|
async def update_run_progress(
|
|
self,
|
|
run_id: str,
|
|
*,
|
|
total_input_tokens: int | None = None,
|
|
total_output_tokens: int | None = None,
|
|
total_tokens: int | None = None,
|
|
llm_call_count: int | None = None,
|
|
lead_agent_tokens: int | None = None,
|
|
subagent_tokens: int | None = None,
|
|
middleware_tokens: int | None = None,
|
|
token_usage_by_model: dict[str, dict[str, int]] | None = None,
|
|
message_count: int | None = None,
|
|
last_ai_message: str | None = None,
|
|
first_human_message: str | None = None,
|
|
) -> None:
|
|
"""Persist a best-effort running snapshot without changing run status."""
|
|
return None
|
|
|
|
@abc.abstractmethod
|
|
async def list_pending(self, *, before: str | None = None) -> list[dict[str, Any]]:
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
async def list_inflight(self, *, before: str | None = None) -> list[dict[str, Any]]:
|
|
"""Return persisted runs that are still ``pending`` or ``running``."""
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
async def aggregate_tokens_by_thread(self, thread_id: str, *, include_active: bool = False) -> dict[str, Any]:
|
|
"""Aggregate token usage for completed runs in a thread.
|
|
|
|
Returns a dict with keys: total_tokens, total_input_tokens,
|
|
total_output_tokens, total_runs, by_model (model_name → {tokens, runs}),
|
|
by_caller ({lead_agent, subagent, middleware}).
|
|
"""
|
|
pass
|