"""Repository layer for database access.""" from typing import Optional from sqlmodel import Session, select from agenteval.models import Case, EvalResult, EvalRun, EvalTarget, Scenario from agenteval.storage.db import ( EvalResultDB, EvalRunDB, EvalTargetDB, ScenarioDB, TurnDB, get_session, utc_now, ) def _target_to_db(target: EvalTarget) -> EvalTargetDB: db = EvalTargetDB( id=target.id, name=target.name, description=target.description, platform=target.platform.value, channel_type=target.channel_type.value, status=target.status.value, created_at=target.created_at, updated_at=target.updated_at or utc_now(), ) db.set_config(target.channel_config) return db def _target_from_db(db: EvalTargetDB) -> EvalTarget: return EvalTarget( id=db.id, name=db.name, description=db.description, platform=db.platform, channel_type=db.channel_type, channel_config=db.get_config(), status=db.status, created_at=db.created_at, updated_at=db.updated_at, ) def _scenario_to_db(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) db.set_cases([case.model_dump() for case in scenario.cases]) db.set_llm_config(scenario.llm_config) return db def _scenario_from_db(db: ScenarioDB) -> Scenario: return Scenario( id=db.id, name=db.name, description=db.description, tags=db.get_tags(), cases=[Case(**case) for case in db.get_cases()], llm_config=db.get_llm_config(), created_at=db.created_at, updated_at=db.updated_at, ) def _run_to_db(run: EvalRun) -> EvalRunDB: db = EvalRunDB( id=run.id, target_id=run.target_id, scenario_id=run.scenario_id, status=run.status.value, started_at=run.started_at, completed_at=run.completed_at, ) if run.summary: db.set_summary(run.summary) return db def _run_from_db(db: EvalRunDB) -> EvalRun: return EvalRun( id=db.id, target_id=db.target_id, scenario_id=db.scenario_id, status=db.status, started_at=db.started_at, completed_at=db.completed_at, summary=db.get_summary(), ) def _result_to_db(result: EvalResult) -> EvalResultDB: return EvalResultDB( id=result.id, run_id=result.run_id, case_id=result.case_id, turn_id=result.turn_id, rule_type=result.rule_type, passed=result.passed, score=result.score, reason=result.reason, ) def _result_from_db(db: EvalResultDB) -> EvalResult: return EvalResult( id=db.id, run_id=db.run_id, case_id=db.case_id, turn_id=db.turn_id, rule_type=db.rule_type, passed=db.passed, score=db.score, reason=db.reason, ) class TargetRepository: """Repository for evaluation targets.""" def __init__(self, session: Optional[Session] = None): self.session = session or get_session() def list_all(self) -> list[EvalTarget]: statement = select(EvalTargetDB).order_by(EvalTargetDB.created_at.desc()) return [_target_from_db(r) for r in self.session.exec(statement).all()] def get(self, target_id: str) -> Optional[EvalTarget]: db = self.session.get(EvalTargetDB, target_id) return _target_from_db(db) if db else None def create(self, target: EvalTarget) -> EvalTarget: db = _target_to_db(target) self.session.add(db) self.session.commit() self.session.refresh(db) return _target_from_db(db) def update(self, target: EvalTarget) -> Optional[EvalTarget]: existing = self.session.get(EvalTargetDB, target.id) if not existing: return None existing.name = target.name existing.description = target.description existing.platform = target.platform.value existing.channel_type = target.channel_type.value existing.status = target.status.value existing.set_config(target.channel_config) existing.updated_at = utc_now() self.session.add(existing) self.session.commit() self.session.refresh(existing) return _target_from_db(existing) def delete(self, target_id: str) -> bool: db = self.session.get(EvalTargetDB, target_id) if not db: return False self.session.delete(db) self.session.commit() return True class ScenarioRepository: """Repository for evaluation scenarios.""" def __init__(self, session: Optional[Session] = None): self.session = session or get_session() def list_all(self) -> list[Scenario]: statement = select(ScenarioDB).order_by(ScenarioDB.created_at.desc()) return [_scenario_from_db(r) for r in self.session.exec(statement).all()] def get(self, scenario_id: str) -> Optional[Scenario]: db = self.session.get(ScenarioDB, scenario_id) return _scenario_from_db(db) if db else None def create(self, scenario: Scenario) -> Scenario: db = _scenario_to_db(scenario) self.session.add(db) self.session.commit() self.session.refresh(db) return _scenario_from_db(db) def update(self, scenario: Scenario) -> Optional[Scenario]: existing = self.session.get(ScenarioDB, scenario.id) if not existing: return None existing.name = scenario.name existing.description = scenario.description existing.set_tags(scenario.tags) existing.set_cases([case.model_dump() for case in scenario.cases]) existing.set_llm_config(scenario.llm_config) existing.updated_at = utc_now() self.session.add(existing) self.session.commit() self.session.refresh(existing) return _scenario_from_db(existing) def delete(self, scenario_id: str) -> bool: db = self.session.get(ScenarioDB, scenario_id) if not db: return False self.session.delete(db) self.session.commit() return True class RunRepository: """Repository for evaluation runs.""" def __init__(self, session: Optional[Session] = None): self.session = session or get_session() def list_all(self) -> list[EvalRun]: statement = select(EvalRunDB).order_by(EvalRunDB.started_at.desc()) return [_run_from_db(r) for r in self.session.exec(statement).all()] def get(self, run_id: str) -> Optional[EvalRun]: db = self.session.get(EvalRunDB, run_id) return _run_from_db(db) if db else None def create(self, run: EvalRun) -> EvalRun: db = _run_to_db(run) self.session.add(db) self.session.commit() self.session.refresh(db) return _run_from_db(db) def update(self, run: EvalRun) -> Optional[EvalRun]: existing = self.session.get(EvalRunDB, run.id) if not existing: return None existing.target_id = run.target_id existing.scenario_id = run.scenario_id existing.status = run.status.value existing.completed_at = run.completed_at if run.summary: existing.set_summary(run.summary) self.session.add(existing) self.session.commit() self.session.refresh(existing) return _run_from_db(existing) def get_turns(self, run_id: str) -> list[TurnDB]: statement = select(TurnDB).where(TurnDB.run_id == run_id).order_by(TurnDB.sent_at) return list(self.session.exec(statement).all()) def get_results(self, run_id: str) -> list[EvalResult]: statement = select(EvalResultDB).where(EvalResultDB.run_id == run_id) return [_result_from_db(r) for r in self.session.exec(statement).all()] def delete(self, run_id: str) -> bool: """Delete a run. ORM-level cascade removes associated turns/results.""" db = self.session.get(EvalRunDB, run_id) if not db: return False self.session.delete(db) self.session.commit() return True class ResultRepository: """Repository for evaluation results.""" def __init__(self, session: Optional[Session] = None): self.session = session or get_session() def save_turn(self, turn) -> TurnDB: db = TurnDB( id=turn.id, run_id=turn.run_id, case_id=turn.case_id, round_index=turn.round_index, question_msg_id=turn.question_msg_id, sent_at=turn.sent_at, received_at=turn.received_at, latency_ms=turn.latency_ms, ) db.set_sent_message(turn.sent_message) db.set_reply(turn.reply) self.session.add(db) self.session.commit() self.session.refresh(db) return db def save_result(self, result: EvalResult) -> EvalResult: db = _result_to_db(result) self.session.add(db) self.session.commit() self.session.refresh(db) return _result_from_db(db)