mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-08-15 09:19:02 +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.
52 lines
2.3 KiB
Python
52 lines
2.3 KiB
Python
"""ORM model for run metadata."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from datetime import UTC, datetime
|
|
|
|
from sqlalchemy import JSON, DateTime, Index, String, Text
|
|
from sqlalchemy.orm import Mapped, mapped_column
|
|
|
|
from deerflow.persistence.base import Base
|
|
|
|
|
|
class RunRow(Base):
|
|
__tablename__ = "runs"
|
|
|
|
run_id: Mapped[str] = mapped_column(String(64), primary_key=True)
|
|
thread_id: Mapped[str] = mapped_column(String(64), nullable=False, index=True)
|
|
assistant_id: Mapped[str | None] = mapped_column(String(128))
|
|
user_id: Mapped[str | None] = mapped_column(String(64), index=True)
|
|
status: Mapped[str] = mapped_column(String(20), default="pending")
|
|
# "pending" | "running" | "success" | "error" | "timeout" | "interrupted"
|
|
|
|
model_name: Mapped[str | None] = mapped_column(String(128))
|
|
multitask_strategy: Mapped[str] = mapped_column(String(20), default="reject")
|
|
metadata_json: Mapped[dict] = mapped_column(JSON, default=dict)
|
|
kwargs_json: Mapped[dict] = mapped_column(JSON, default=dict)
|
|
error: Mapped[str | None] = mapped_column(Text)
|
|
|
|
# Convenience fields (for listing pages without querying RunEventStore)
|
|
message_count: Mapped[int] = mapped_column(default=0)
|
|
first_human_message: Mapped[str | None] = mapped_column(Text)
|
|
last_ai_message: Mapped[str | None] = mapped_column(Text)
|
|
|
|
# Token usage (accumulated in-memory by RunJournal, written on run completion)
|
|
total_input_tokens: Mapped[int] = mapped_column(default=0)
|
|
total_output_tokens: Mapped[int] = mapped_column(default=0)
|
|
total_tokens: Mapped[int] = mapped_column(default=0)
|
|
llm_call_count: Mapped[int] = mapped_column(default=0)
|
|
lead_agent_tokens: Mapped[int] = mapped_column(default=0)
|
|
subagent_tokens: Mapped[int] = mapped_column(default=0)
|
|
middleware_tokens: Mapped[int] = mapped_column(default=0)
|
|
# Per-model token breakdown
|
|
token_usage_by_model: Mapped[dict] = mapped_column(JSON, default=dict)
|
|
|
|
# Follow-up association
|
|
follow_up_to_run_id: Mapped[str | None] = mapped_column(String(64))
|
|
|
|
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC))
|
|
updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC), onupdate=lambda: datetime.now(UTC))
|
|
|
|
__table_args__ = (Index("ix_runs_thread_status", "thread_id", "status"),)
|