AgentEvalTool/backend/agenteval/storage/repository/scenario.py

"""Repository for evaluation scenarios."""

from typing import Optional

from sqlmodel import select

from agenteval.models import Case, Scenario
from agenteval.services.model_configs import ModelConfigService
from agenteval.storage.db import ScenarioDB, utc_now
from agenteval.storage.model_config_repository import ScenarioModelBindingRepository
from agenteval.storage.repository.base import BaseRepository


class ScenarioRepository(BaseRepository[Scenario, ScenarioDB]):
    """Repository for evaluation scenarios."""

    _table = ScenarioDB
    _order_by = "created_at"

    def _to_db(self, scenario: Scenario) -> ScenarioDB:
        db = ScenarioDB(
            id=scenario.id,
            name=scenario.name,
            description=scenario.description,
            created_at=scenario.created_at,
            updated_at=scenario.updated_at or utc_now(),
        )
        db.set_tags(scenario.tags)
        # mode="json" 与 update() 的考纲比较保持同一序列化形态，避免假升版
        db.set_cases([case.model_dump(mode="json") for case in scenario.cases])
        db.set_llm_config(scenario.llm_config)
        db.set_acceptance_criteria(scenario.acceptance_criteria)
        return db

    def _from_db(self, db: ScenarioDB) -> Scenario:
        bindings = ScenarioModelBindingRepository(self.session).get_for_scenario(db.id or "")
        return Scenario(
            id=db.id,
            name=db.name,
            description=db.description,
            tags=db.get_tags(),
            cases=[Case(**case) for case in db.get_cases()],
            model_bindings=bindings,
            llm_config=db.get_llm_config(),
            acceptance_criteria=db.get_acceptance_criteria(),
            version=db.version or 1,
            created_at=db.created_at,
            updated_at=db.updated_at,
        )

    def create(self, scenario: Scenario) -> Scenario:
        db = self._to_db(scenario)
        bindings = {purpose.value: config_id for purpose, config_id in scenario.model_bindings.items()}
        try:
            ModelConfigService(self.session).validate_bindings(bindings)
            self.session.add(db)
            self.session.flush()
            ScenarioModelBindingRepository(self.session).replace_for_scenario(db.id or "", bindings)
            self.session.commit()
            self.session.refresh(db)
        except Exception:
            self.session.rollback()
            raise
        return self._from_db(db)

    def update(self, scenario: Scenario) -> Optional[Scenario]:
        existing = self.session.get(ScenarioDB, scenario.id)
        if not existing:
            return None
        bindings = {purpose.value: config_id for purpose, config_id in scenario.model_bindings.items()}
        try:
            ModelConfigService(self.session).validate_bindings(bindings)
            # 考纲字段（cases / model_bindings / llm_config）变更才升版（ADR-0001）；
            # 版本由系统维护，忽略 scenario.version 的外部传入值。
            new_cases = [case.model_dump(mode="json") for case in scenario.cases]
            old_bindings = ScenarioModelBindingRepository(self.session).get_for_scenario(existing.id or "")
            syllabus_changed = (
                existing.get_cases() != new_cases
                or existing.get_llm_config() != scenario.llm_config
                or old_bindings != bindings
            )
            if syllabus_changed:
                existing.version = (existing.version or 1) + 1
            existing.name = scenario.name
            existing.description = scenario.description
            existing.set_tags(scenario.tags)
            existing.set_cases(new_cases)
            existing.set_llm_config(scenario.llm_config)
            existing.set_acceptance_criteria(scenario.acceptance_criteria)
            existing.updated_at = utc_now()
            self.session.add(existing)
            ScenarioModelBindingRepository(self.session).replace_for_scenario(existing.id or "", bindings)
            self.session.commit()
            self.session.refresh(existing)
        except Exception:
            self.session.rollback()
            raise
        return self._from_db(existing)

    def delete(self, scenario_id: str) -> bool:
        db = self.session.get(ScenarioDB, scenario_id)
        if not db:
            return False
        try:
            ScenarioModelBindingRepository(self.session).delete_for_scenario(scenario_id)
            self.session.delete(db)
            self.session.commit()
        except Exception:
            self.session.rollback()
            raise
        return True

    def list_all(self) -> list[Scenario]:
        """Batch load scenarios with bindings to avoid N+1 queries."""
        column = getattr(self._table, self._order_by)
        statement = select(self._table).order_by(column.desc())
        scenarios = self.session.exec(statement).all()

        # Batch load all bindings at once
        all_bindings = ScenarioModelBindingRepository(self.session).get_all_for_scenarios(
            [s.id for s in scenarios if s.id]
        )

        return [self._from_db_with_bindings(s, all_bindings.get(s.id or "", [])) for s in scenarios]

    def _from_db_with_bindings(self, db: ScenarioDB, bindings: list) -> Scenario:
        """Convert DB model to domain model with pre-loaded bindings."""
        return Scenario(
            id=db.id,
            name=db.name,
            description=db.description,
            tags=db.get_tags(),
            cases=[Case(**case) for case in db.get_cases()],
            model_bindings=bindings,
            llm_config=db.get_llm_config(),
            acceptance_criteria=db.get_acceptance_criteria(),
            version=db.version or 1,
            created_at=db.created_at,
            updated_at=db.updated_at,
        )

    def name_map(self) -> dict[str, str]:
        """scenario_id → 名称映射：报告 / 时间线 / 列表等读路径共用的场景名取法。"""
        return {sid: name for sid, name in self.session.exec(select(ScenarioDB.id, ScenarioDB.name)).all()}