aa40c8e0d8
Some checks failed
CI / test (push) Failing after 12s
Campaign.summary was a bare Optional[dict] while RunSummary is a typed VO — scheduler state (spawned_indices/errors) flowed untyped through set_/get_summary. Introduce CampaignSummary + SchedulerState (extra=allow, validate_assignment), mirroring RunSummary; campaign_runner reads/writes the VO. Also converge the ~10 repeated JSON column get/set pairs onto _json_dumps/_json_loads helpers, unifying ensure_ascii=False and fixing the set_modalities ensure_ascii=True trap.
348 lines
11 KiB
Python
348 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
|
|
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
|
|
|
|
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()
|