"""Repository for intelligent evaluation entities.""" from typing import Optional from sqlmodel import Session, select from agenteval.intelligent_eval.models import ( IntelligentEval, IntelligentEvalMessage, IntelligentEvalSession, IntelligentEvalSessionStatus, IntelligentEvalStatus, ) from agenteval.storage.db import ( IntelligentEvalDB, IntelligentEvalMessageDB, IntelligentEvalSessionDB, get_session, utc_now, ) class IntelligentEvalRepository: """CRUD + lifecycle operations for intelligent evaluations.""" def __init__(self, session: Optional[Session] = None): self.session = session or get_session() def _to_db(self, obj: IntelligentEval) -> IntelligentEvalDB: db = IntelligentEvalDB( id=obj.id, name=obj.name, target_id=obj.target_id, status=obj.status.value if isinstance(obj.status, IntelligentEvalStatus) else obj.status, goal=obj.goal, intent=obj.intent, role_description=obj.role_description, time_window_hours=obj.time_window_hours, plan_feedback=obj.plan_feedback, created_at=obj.created_at, updated_at=obj.updated_at, started_at=obj.started_at, completed_at=obj.completed_at, ) db.set_seeds(obj.seeds) if obj.plan is not None: db.set_plan(obj.plan) if obj.report is not None: db.set_report(obj.report) return db def _from_db(self, db: IntelligentEvalDB) -> IntelligentEval: return IntelligentEval( id=db.id, name=db.name, target_id=db.target_id, status=IntelligentEvalStatus(db.status), goal=db.goal, seeds=db.get_seeds(), intent=db.intent, role_description=db.role_description, plan=db.get_plan(), plan_feedback=db.plan_feedback, time_window_hours=db.time_window_hours, report=db.get_report(), created_at=db.created_at, updated_at=db.updated_at, started_at=db.started_at, completed_at=db.completed_at, ) def list_all(self) -> list[IntelligentEval]: statement = select(IntelligentEvalDB).order_by(IntelligentEvalDB.created_at.desc()) return [self._from_db(r) for r in self.session.exec(statement).all()] def get(self, eval_id: str) -> Optional[IntelligentEval]: db = self.session.get(IntelligentEvalDB, eval_id) return self._from_db(db) if db else None def create(self, obj: IntelligentEval) -> IntelligentEval: db = self._to_db(obj) self.session.add(db) self.session.commit() self.session.refresh(db) return self._from_db(db) def update(self, obj: IntelligentEval) -> IntelligentEval: db = self.session.get(IntelligentEvalDB, obj.id) if not db: raise ValueError(f"IntelligentEval {obj.id} not found") updated = self._to_db(obj) updated.id = db.id # Preserve fields that _to_db doesn't set from None self.session.delete(db) self.session.add(updated) self.session.commit() self.session.refresh(updated) return self._from_db(updated) def transition_status(self, eval_id: str, new_status: IntelligentEvalStatus) -> Optional[IntelligentEval]: db = self.session.get(IntelligentEvalDB, eval_id) if not db: return None db.status = new_status.value db.updated_at = utc_now() if new_status == IntelligentEvalStatus.EXECUTING and db.started_at is None: db.started_at = utc_now() if new_status in (IntelligentEvalStatus.COMPLETED, IntelligentEvalStatus.CANCELLED, IntelligentEvalStatus.FAILED): db.completed_at = utc_now() self.session.add(db) self.session.commit() self.session.refresh(db) return self._from_db(db) def delete(self, eval_id: str) -> bool: db = self.session.get(IntelligentEvalDB, eval_id) if not db: return False self.session.delete(db) self.session.commit() return True class IntelligentEvalSessionRepository: """CRUD for intelligent eval sessions.""" def __init__(self, session: Optional[Session] = None): self.session = session or get_session() def _from_db(self, db: IntelligentEvalSessionDB) -> IntelligentEvalSession: return IntelligentEvalSession( id=db.id, eval_id=db.eval_id, target_id=db.target_id, persona=db.get_persona(), goal=db.goal, dimension=db.dimension, status=IntelligentEvalSessionStatus(db.status), verdict=db.get_verdict(), turn_count=db.turn_count, created_at=db.created_at, closed_at=db.closed_at, ) def list_by_eval(self, eval_id: str) -> list[IntelligentEvalSession]: statement = ( select(IntelligentEvalSessionDB) .where(IntelligentEvalSessionDB.eval_id == eval_id) .order_by(IntelligentEvalSessionDB.created_at.asc()) ) return [self._from_db(r) for r in self.session.exec(statement).all()] def get(self, session_id: str) -> Optional[IntelligentEvalSession]: db = self.session.get(IntelligentEvalSessionDB, session_id) return self._from_db(db) if db else None def create(self, obj: IntelligentEvalSession) -> IntelligentEvalSession: db = IntelligentEvalSessionDB( id=obj.id, eval_id=obj.eval_id, target_id=obj.target_id, goal=obj.goal, dimension=obj.dimension, status=obj.status.value if isinstance(obj.status, IntelligentEvalSessionStatus) else obj.status, turn_count=obj.turn_count, ) db.set_persona(obj.persona) if obj.verdict is not None: db.set_verdict(obj.verdict) self.session.add(db) self.session.commit() self.session.refresh(db) return self._from_db(db) def close(self, session_id: str, verdict: dict, status: IntelligentEvalSessionStatus = IntelligentEvalSessionStatus.COMPLETED) -> Optional[IntelligentEvalSession]: db = self.session.get(IntelligentEvalSessionDB, session_id) if not db: return None db.status = status.value db.set_verdict(verdict) db.closed_at = utc_now() self.session.add(db) self.session.commit() self.session.refresh(db) return self._from_db(db) def increment_turns(self, session_id: str) -> None: db = self.session.get(IntelligentEvalSessionDB, session_id) if db: db.turn_count += 1 self.session.add(db) self.session.commit() class IntelligentEvalMessageRepository: """CRUD for intelligent eval session messages.""" def __init__(self, session: Optional[Session] = None): self.session = session or get_session() def _from_db(self, db: IntelligentEvalMessageDB) -> IntelligentEvalMessage: return IntelligentEvalMessage( id=db.id, session_id=db.session_id, role=db.role, content=db.content, latency_ms=db.latency_ms, created_at=db.created_at, ) def list_by_session(self, session_id: str) -> list[IntelligentEvalMessage]: statement = ( select(IntelligentEvalMessageDB) .where(IntelligentEvalMessageDB.session_id == session_id) .order_by(IntelligentEvalMessageDB.created_at.asc()) ) return [self._from_db(r) for r in self.session.exec(statement).all()] def create(self, obj: IntelligentEvalMessage) -> IntelligentEvalMessage: db = IntelligentEvalMessageDB( id=obj.id, session_id=obj.session_id, role=obj.role, content=obj.content, latency_ms=obj.latency_ms, created_at=obj.created_at, ) self.session.add(db) self.session.commit() self.session.refresh(db) return self._from_db(db)