"""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 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()