c63236f671
Campaigns can pin an analysis model config instead of following the global analysis default. Creation validates the referenced config exists (400 otherwise); the create form offers enabled chat configs with the global default as the fallback option.
350 lines
11 KiB
Python
350 lines
11 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 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()
|