AgentEvalTool/backend/agenteval/intelligent_eval/repository.py

"""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)