AgentEvalTool/backend/agenteval/storage/db.py
sinohqb 739d586aec feat(backend): v0.4 triggered_by tracking, login gate, compare guard, dashboard stats
- EvalRun.triggered_by 全链路(manual/ai_assistant/cli)+ 迁移 b7d4e6f81c22
- 标准 agenteval-run SKILL.md 纳入版本管理,deploy 脚本同步 + API Key 注入
- 简单登录:AGENTEVAL_ADMIN_PASSWORD + HMAC 会话 token,require_auth 双凭据
- 对比报告限同场景(400)+ 空 results 误判修复
- /api/stats/dashboard 扩展聚合;/api/runs 返回场景/对象名
- 测试 218 → 232
2026-07-28 17:40:54 +08:00

305 lines
9.8 KiB
Python

"""Database models and session management using SQLModel."""
import json
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Optional
from sqlalchemy.pool import StaticPool
from sqlmodel import Field, Relationship, Session, SQLModel, create_engine
DATA_DIR = Path(__file__).resolve().parent.parent.parent.parent / "data"
DATA_DIR.mkdir(parents=True, exist_ok=True)
DATABASE_URL = f"sqlite:///{DATA_DIR / 'agenteval.db'}"
FILES_DIR = DATA_DIR / "files"
FILES_DIR.mkdir(parents=True, exist_ok=True)
engine = create_engine(
DATABASE_URL,
echo=False,
connect_args={"check_same_thread": False},
poolclass=StaticPool,
)
def utc_now() -> datetime:
return datetime.now(timezone.utc)
def iso_utc(dt: datetime | None) -> str | None:
"""Serialize a datetime to ISO 8601 with UTC timezone suffix.
Guarantees the output always ends with 'Z' or '+00:00' so JavaScript's
Date.parse() interprets it correctly as UTC (no 8-hour local-time offset).
"""
if dt is None:
return None
if dt.tzinfo is None:
dt = dt.replace(tzinfo=timezone.utc)
return dt.isoformat(timespec="seconds").replace("+00:00", "Z")
def new_uuid() -> str:
return str(uuid.uuid4())
class EvalTargetDB(SQLModel, table=True):
"""Database table for evaluation targets."""
__tablename__ = "eval_targets"
id: Optional[str] = Field(default_factory=new_uuid, primary_key=True)
name: str
description: str = ""
platform: str = "ai_digital_employee"
channel_type: str = "tutu-api"
channel_config: str = "{}"
status: str = "active"
created_at: Optional[datetime] = Field(default_factory=utc_now)
updated_at: Optional[datetime] = Field(default_factory=utc_now)
runs: list["EvalRunDB"] = Relationship(
back_populates="target",
sa_relationship_kwargs={"cascade": "all, delete-orphan"},
)
def get_config(self) -> dict[str, Any]:
return json.loads(self.channel_config)
def set_config(self, config: dict[str, Any]) -> None:
self.channel_config = json.dumps(config, ensure_ascii=False)
class ScenarioDB(SQLModel, table=True):
"""Database table for evaluation scenarios."""
__tablename__ = "scenarios"
id: Optional[str] = Field(default_factory=new_uuid, primary_key=True)
name: str
description: str = ""
tags: str = "[]"
cases: str = "[]"
llm_config: Optional[str] = None
created_at: Optional[datetime] = Field(default_factory=utc_now)
updated_at: Optional[datetime] = Field(default_factory=utc_now)
runs: list["EvalRunDB"] = Relationship(
back_populates="scenario",
sa_relationship_kwargs={"cascade": "all, delete-orphan"},
)
def get_tags(self) -> list[str]:
return json.loads(self.tags)
def set_tags(self, tags: list[str]) -> None:
self.tags = json.dumps(tags, ensure_ascii=False)
def get_cases(self) -> list[dict[str, Any]]:
return json.loads(self.cases)
def set_cases(self, cases: list[dict[str, Any]]) -> None:
self.cases = json.dumps(cases, ensure_ascii=False)
def get_llm_config(self) -> Optional[dict[str, Any]]:
return json.loads(self.llm_config) if self.llm_config else None
def set_llm_config(self, config: Optional[dict[str, Any]]) -> None:
self.llm_config = json.dumps(config, ensure_ascii=False) if config else None
class ModelConfigDB(SQLModel, table=True):
"""Reusable external model connection configuration."""
__tablename__ = "model_configs"
id: Optional[str] = Field(default_factory=new_uuid, primary_key=True)
name: str = Field(index=True, unique=True)
provider: str = "openai_compatible"
capability: str = Field(index=True)
endpoint_url: str
model_name: Optional[str] = None
vendor_name: str = ""
input_modalities: str = '["text"]'
output_modalities: str = '["text"]'
context_window: Optional[int] = None
max_output_tokens: Optional[int] = None
supports_streaming: bool = False
supports_tool_calling: bool = False
supports_structured_output: bool = False
supports_reasoning: bool = False
region: str = ""
documentation_url: Optional[str] = None
api_key_encrypted: Optional[str] = None
enabled: bool = True
is_default: bool = False
description: str = ""
created_at: Optional[datetime] = Field(default_factory=utc_now)
updated_at: Optional[datetime] = Field(default_factory=utc_now)
def get_input_modalities(self) -> list[str]:
return json.loads(self.input_modalities)
def get_output_modalities(self) -> list[str]:
return json.loads(self.output_modalities)
def set_modalities(self, input_modalities: list[str], output_modalities: list[str]) -> None:
self.input_modalities = json.dumps(input_modalities, ensure_ascii=True)
self.output_modalities = json.dumps(output_modalities, ensure_ascii=True)
class ScenarioModelBindingDB(SQLModel, table=True):
"""Bind one model configuration to a purpose within a scenario."""
__tablename__ = "scenario_model_bindings"
scenario_id: str = Field(foreign_key="scenarios.id", primary_key=True)
purpose: str = Field(primary_key=True)
model_config_id: str = Field(foreign_key="model_configs.id", index=True)
class EvalRunDB(SQLModel, table=True):
"""Database table for evaluation runs."""
__tablename__ = "eval_runs"
id: Optional[str] = Field(default_factory=new_uuid, primary_key=True)
target_id: Optional[str] = Field(default=None, foreign_key="eval_targets.id")
scenario_id: Optional[str] = Field(default=None, foreign_key="scenarios.id")
status: str = "pending"
triggered_by: str = Field(default="manual")
started_at: Optional[datetime] = Field(default_factory=utc_now)
completed_at: Optional[datetime] = None
summary: Optional[str] = None
target: Optional[EvalTargetDB] = Relationship(back_populates="runs")
scenario: Optional[ScenarioDB] = Relationship(back_populates="runs")
turns: list["TurnDB"] = Relationship(
back_populates="run",
sa_relationship_kwargs={"cascade": "all, delete-orphan"},
)
results: list["EvalResultDB"] = Relationship(
back_populates="run",
sa_relationship_kwargs={"cascade": "all, delete-orphan"},
)
def get_summary(self) -> Optional[dict[str, Any]]:
return json.loads(self.summary) if self.summary else None
def set_summary(self, summary: dict[str, Any]) -> None:
self.summary = json.dumps(summary, ensure_ascii=False)
class TurnDB(SQLModel, table=True):
"""Database table for conversation turns."""
__tablename__ = "turns"
id: Optional[str] = Field(default_factory=new_uuid, primary_key=True)
run_id: Optional[str] = Field(default=None, foreign_key="eval_runs.id")
case_id: str
round_index: int
sent_message: str = "{}"
sent_at: Optional[datetime] = Field(default_factory=utc_now)
question_msg_id: Optional[str] = None
reply: Optional[str] = None
received_at: Optional[datetime] = None
latency_ms: Optional[int] = None
run: Optional[EvalRunDB] = Relationship(back_populates="turns")
def get_sent_message(self) -> dict[str, Any]:
return json.loads(self.sent_message)
def set_sent_message(self, message: dict[str, Any]) -> None:
self.sent_message = json.dumps(message, ensure_ascii=False)
def get_reply(self) -> Optional[dict[str, Any]]:
return json.loads(self.reply) if self.reply else None
def set_reply(self, reply: Optional[dict[str, Any]]) -> None:
self.reply = json.dumps(reply, ensure_ascii=False) if reply else None
class EvalResultDB(SQLModel, table=True):
"""Database table for evaluation results."""
__tablename__ = "eval_results"
id: Optional[str] = Field(default_factory=new_uuid, primary_key=True)
run_id: Optional[str] = Field(default=None, foreign_key="eval_runs.id")
case_id: str
turn_id: str
rule_type: str
passed: bool
score: Optional[float] = None
reason: str = ""
run: Optional[EvalRunDB] = Relationship(back_populates="results")
class FileCategoryDB(SQLModel, table=True):
"""Database table for file categories (tree structure via parent_id)."""
__tablename__ = "file_categories"
id: Optional[str] = Field(default_factory=new_uuid, primary_key=True)
name: str
parent_id: Optional[str] = Field(
default=None,
foreign_key="file_categories.id",
)
created_at: Optional[datetime] = Field(default_factory=utc_now)
updated_at: Optional[datetime] = Field(default_factory=utc_now)
parent: Optional["FileCategoryDB"] = Relationship(
back_populates="children",
sa_relationship_kwargs={"remote_side": "FileCategoryDB.id"},
)
children: list["FileCategoryDB"] = Relationship(
back_populates="parent",
sa_relationship_kwargs={"cascade": "all, delete-orphan"},
)
files: list["FileRecordDB"] = Relationship(
back_populates="category",
sa_relationship_kwargs={"cascade": "all, delete-orphan"},
)
class FileRecordDB(SQLModel, table=True):
"""Database table for uploaded file records."""
__tablename__ = "file_records"
id: Optional[str] = Field(default_factory=new_uuid, primary_key=True)
original_name: str
storage_name: str
category_id: Optional[str] = Field(
default=None,
foreign_key="file_categories.id",
)
file_size: int = 0
mime_type: str = ""
file_ext: str = ""
created_at: Optional[datetime] = Field(default_factory=utc_now)
category: Optional[FileCategoryDB] = Relationship(back_populates="files")
def init_db() -> None:
SQLModel.metadata.create_all(engine)
def get_session() -> Session:
return Session(engine)
def get_session_context():
"""Context manager that creates and properly closes a database session."""
session = Session(engine)
try:
yield session
finally:
session.close()