a351f65550
v0.9 ticket 01. Independent exploration_sessions/exploration_messages entities (never merged into EvalRun, keeping ADR-0001/0002 semantics intact): create/message/close APIs forward virtual-user messages through the target's real channel, persist both parties' rows with latency, and close with a whitelist-normalized experience record. Budget enforcement is a platform ledger — sessions per window, turns per session, and session interval overruns return 409 with readable reasons; accelerated lines accept manual sessions only. Messages delivered but unanswered still consume a turn so timeouts cannot bypass the budget.
468 lines
16 KiB
Python
468 lines
16 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 CampaignPeriodComparisonDB(SQLModel, table=True):
|
|
"""One row per campaign holding its period-comparison (周期对比) state.
|
|
|
|
Mirrors ``campaign_analyses`` (one row, upserted per campaign), plus the
|
|
baseline campaign this comparison was computed against.
|
|
"""
|
|
|
|
__tablename__ = "campaign_period_comparisons"
|
|
|
|
id: Optional[str] = Field(default_factory=new_uuid, primary_key=True)
|
|
campaign_id: str = Field(unique=True, index=True)
|
|
baseline_campaign_id: str = Field(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 ExplorationSessionDB(SQLModel, table=True):
|
|
"""One virtual-user exploration session (探索会话) within a campaign.
|
|
|
|
Independent entity, deliberately NOT an EvalRun: exploration outcomes feed
|
|
the 体验判定 evidence line and must not pollute pass-rate semantics
|
|
(ADR-0001 comparability / ADR-0002).
|
|
"""
|
|
|
|
__tablename__ = "exploration_sessions"
|
|
|
|
id: Optional[str] = Field(default_factory=new_uuid, primary_key=True)
|
|
campaign_id: str = Field(index=True, foreign_key="campaigns.id")
|
|
target_id: str = Field(foreign_key="eval_targets.id")
|
|
persona: str = "{}"
|
|
goal: str = ""
|
|
seed_ref: Optional[str] = None
|
|
status: str = "running"
|
|
triggered_by: str = "auto"
|
|
experience: Optional[str] = None
|
|
judge_review: Optional[str] = None
|
|
turn_count: int = 0
|
|
error: Optional[str] = None
|
|
created_at: Optional[datetime] = Field(default_factory=utc_now)
|
|
closed_at: Optional[datetime] = None
|
|
|
|
def get_persona(self) -> dict[str, Any]:
|
|
return _json_loads(self.persona)
|
|
|
|
def set_persona(self, persona: dict[str, Any]) -> None:
|
|
self.persona = _json_dumps(persona)
|
|
|
|
def get_seed_ref(self) -> Optional[dict[str, Any]]:
|
|
return _json_loads(self.seed_ref) if self.seed_ref else None
|
|
|
|
def set_seed_ref(self, seed_ref: dict[str, Any]) -> None:
|
|
self.seed_ref = _json_dumps(seed_ref)
|
|
|
|
def get_experience(self) -> Optional[dict[str, Any]]:
|
|
return _json_loads(self.experience) if self.experience else None
|
|
|
|
def set_experience(self, experience: dict[str, Any]) -> None:
|
|
self.experience = _json_dumps(experience)
|
|
|
|
def get_judge_review(self) -> Optional[dict[str, Any]]:
|
|
return _json_loads(self.judge_review) if self.judge_review else None
|
|
|
|
def set_judge_review(self, judge_review: dict[str, Any]) -> None:
|
|
self.judge_review = _json_dumps(judge_review)
|
|
|
|
|
|
class ExplorationMessageDB(SQLModel, table=True):
|
|
"""One chat message inside an exploration session (role/content form).
|
|
|
|
Parallel to ``eval_runs`` turns but never shared with them. A completed
|
|
question-answer pair contributes two rows (user + assistant); the reply
|
|
row carries ``latency_ms``.
|
|
"""
|
|
|
|
__tablename__ = "exploration_messages"
|
|
|
|
id: Optional[str] = Field(default_factory=new_uuid, primary_key=True)
|
|
session_id: str = Field(index=True, foreign_key="exploration_sessions.id")
|
|
round_index: int = 0
|
|
role: str = "user"
|
|
content: str = ""
|
|
latency_ms: Optional[int] = None
|
|
created_at: Optional[datetime] = Field(default_factory=utc_now)
|
|
|
|
|
|
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()
|