AgentEvalTool/backend/agenteval/storage/db.py
sinohqb 15c542d92c feat(analysis): two-phase campaign analysis agent with storage and API
Add the analysis role's execution path: a two-phase orchestration
(per-scenario diagnosis gathered in parallel, then a synthesis pass)
that reads the existing campaign report aggregation plus capped failure
samples, validates the LLM's JSON against the report schema, and strips
fabricated run/scenario references before persisting. Results upsert one
row per campaign (generating/completed/failed) with the model config
snapshot; GET/POST /api/campaigns/{id}/analysis expose the state machine,
guarding non-terminal campaigns and missing analysis models with 400s.
2026-08-03 02:06:29 +08:00

372 lines
12 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())
def _json_dumps(value: Any) -> str:
"""Serialize a JSON column value. ensure_ascii=False keeps CJK readable
in the stored text — the single serialization口径 for all JSON columns."""
return json.dumps(value, ensure_ascii=False)
def _json_loads(raw: str) -> Any:
return json.loads(raw)
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)
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
version: int = Field(default=1)
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)
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)
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) 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
is_analysis_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)
self.output_modalities = _json_dumps(output_modalities)
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 CampaignDB(SQLModel, table=True):
"""Database table for evaluation campaigns (评估活动)."""
__tablename__ = "campaigns"
id: Optional[str] = Field(default_factory=new_uuid, primary_key=True)
name: str
target_id: Optional[str] = Field(default=None, foreign_key="eval_targets.id")
window_seconds: int
time_scale: float = 1.0
plan: str = "[]"
status: str = "planned"
started_at: Optional[datetime] = None
completed_at: Optional[datetime] = None
created_at: Optional[datetime] = Field(default_factory=utc_now)
summary: Optional[str] = None
analysis_model_config_id: Optional[str] = None
def get_plan(self) -> list[dict[str, Any]]:
return _json_loads(self.plan)
def set_plan(self, plan: list[dict[str, Any]]) -> None:
self.plan = _json_dumps(plan)
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)
class CampaignAnalysisDB(SQLModel, table=True):
"""One row per campaign holding its intelligent analysis (智能分析) state."""
__tablename__ = "campaign_analyses"
id: Optional[str] = Field(default_factory=new_uuid, primary_key=True)
campaign_id: str = Field(unique=True, index=True)
status: str = "generating"
result: Optional[str] = None
model_config_id: Optional[str] = None
error: Optional[str] = None
triggered_by: str = "manual"
created_at: Optional[datetime] = Field(default_factory=utc_now)
updated_at: Optional[datetime] = Field(default_factory=utc_now)
def get_result(self) -> Optional[dict[str, Any]]:
return _json_loads(self.result) if self.result else None
def set_result(self, result: dict[str, Any]) -> None:
self.result = _json_dumps(result)
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")
scenario_version: int = Field(default=1)
campaign_id: Optional[str] = Field(default=None, foreign_key="campaigns.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)
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)
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) 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()