AgentEvalTool/backend/agenteval/storage/db/eval_tables.py

"""Core evaluation tables: targets, scenarios, runs, turns, results."""

from datetime import datetime
from typing import Any, Optional

import sqlalchemy as sa
from sqlmodel import Field, Relationship, SQLModel

from agenteval.storage.db.core import _json_dumps, _json_loads, new_uuid, utc_now


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 EvalRunDB(SQLModel, table=True):
    """Database table for evaluation runs."""

    __tablename__ = "eval_runs"
    __table_args__ = (
        sa.Index(
            "uq_eval_runs_campaign_occurrence",
            "campaign_id",
            "campaign_plan_index",
            "campaign_occurrence_index",
            unique=True,
        ),
        sa.Index("ix_eval_runs_status", "status"),
        sa.Index("ix_eval_runs_campaign_id", "campaign_id"),
    )

    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")
    campaign_plan_index: Optional[int] = Field(default=None, ge=0)
    campaign_occurrence_index: Optional[int] = Field(default=None, ge=0)
    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"
    __table_args__ = (sa.Index("ix_turns_run_id", "run_id"),)

    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"
    __table_args__ = (sa.Index("ix_eval_results_run_id", "run_id"),)

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