"""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) 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(), 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.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(), 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()}