mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-05-23 00:16:48 +00:00
perf: use SQL aggregation for feedback stats and thread token usage
Replace Python-side counting in FeedbackRepository.aggregate_by_run with a single SELECT COUNT/SUM query. Add RunStore.aggregate_tokens_by_thread abstract method with SQL GROUP BY implementation in RunRepository and Python fallback in MemoryRunStore. Simplify the thread_token_usage endpoint to delegate to the new method, eliminating the limit=10000 truncation risk. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -84,3 +84,13 @@ class RunStore(abc.ABC):
|
||||
@abc.abstractmethod
|
||||
async def list_pending(self, *, before: str | None = None) -> list[dict[str, Any]]:
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
async def aggregate_tokens_by_thread(self, thread_id: str) -> 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
|
||||
|
||||
@@ -77,3 +77,24 @@ class MemoryRunStore(RunStore):
|
||||
results = [r for r in self._runs.values() if r["status"] == "pending" and r["created_at"] <= now]
|
||||
results.sort(key=lambda r: r["created_at"])
|
||||
return results
|
||||
|
||||
async def aggregate_tokens_by_thread(self, thread_id: str) -> dict[str, Any]:
|
||||
completed = [r for r in self._runs.values() if r["thread_id"] == thread_id and r.get("status") in ("success", "error")]
|
||||
by_model: dict[str, dict] = {}
|
||||
for r in completed:
|
||||
model = r.get("model_name") or "unknown"
|
||||
entry = by_model.setdefault(model, {"tokens": 0, "runs": 0})
|
||||
entry["tokens"] += r.get("total_tokens", 0)
|
||||
entry["runs"] += 1
|
||||
return {
|
||||
"total_tokens": sum(r.get("total_tokens", 0) for r in completed),
|
||||
"total_input_tokens": sum(r.get("total_input_tokens", 0) for r in completed),
|
||||
"total_output_tokens": sum(r.get("total_output_tokens", 0) for r in completed),
|
||||
"total_runs": len(completed),
|
||||
"by_model": by_model,
|
||||
"by_caller": {
|
||||
"lead_agent": sum(r.get("lead_agent_tokens", 0) for r in completed),
|
||||
"subagent": sum(r.get("subagent_tokens", 0) for r in completed),
|
||||
"middleware": sum(r.get("middleware_tokens", 0) for r in completed),
|
||||
},
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user