mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-05-14 20:53:41 +00:00
de253e4a0a
* feat(run): propagate model_name from gateway request context to persistence layer Pass model_name through the full run creation pipeline — from RunCreateRequest.context in the gateway, through RunManager, to the RunStore interface and SQL persistence. This enables client-specified model selection to be recorded per-run in the database. * feat(run): add model allowlist validation and effective model name capture - Validate model_name against allowlist in gateway services.py using get_app_config().get_model_config() - Truncate model_name to 128 chars to match DB column constraint - In worker.py, capture effective model name from agent.metadata after agent creation and persist if resolved differently than requested * feat(run): add defense-in-depth model_name normalization and round-trip persistence tests - Add _normalize_model_name() to RunRepository for whitespace stripping and 128-char truncation before DB writes. - Add round-trip unit tests for model_name creation and default None in test_run_manager.py. * fix(run): coerce non-string model_name values before strip/truncate in _normalize_model_name * fix(gateway): add runtime type guard for model_name coercion in gateway services Add isinstance check and str() coercion before calling .strip() to prevent AttributeError when non-string types (int, None, etc.) flow through the gateway. Paired with SQL integration test for end-to-end model_name persistence across gateway → langgraph → persistence layer. * fix(run): drop Alembic migration for model_name (no-op) and expose public update method on RunManager - Drop a1b2c3d4e5f6 migration: model_name already exists in RunRow schema and is auto-created via Base.metadata.create_all() at startup - Add update_model_name() public method to RunManager to replace the private _persist_to_store call in worker.py, preserving internal locking/persistence
97 lines
2.6 KiB
Python
97 lines
2.6 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) -> 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,
|
|
) -> None:
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
async def delete(self, run_id: str) -> None:
|
|
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,
|
|
message_count: int = 0,
|
|
last_ai_message: str | None = None,
|
|
first_human_message: str | None = None,
|
|
error: str | None = None,
|
|
) -> None:
|
|
pass
|
|
|
|
@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
|