AgentEvalTool/backend/agenteval/storage/db.py

"""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 new_uuid() -> str:
    return str(uuid.uuid4())


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, ensure_ascii=False)


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
    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, ensure_ascii=False)

    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, ensure_ascii=False)

    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, ensure_ascii=False) if config else None


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")
    status: str = "pending"
    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, ensure_ascii=False)


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, ensure_ascii=False)

    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, ensure_ascii=False) 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()