AgentEvalTool/backend/agenteval/intelligent_eval/repository.py

"""Repository for intelligent evaluation entities."""

import json
from collections.abc import Sequence
from dataclasses import dataclass
from enum import Enum
from typing import Optional

from sqlalchemy import func, literal
from sqlalchemy import insert as sql_insert
from sqlalchemy import update as sql_update
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,
    new_uuid,
    utc_now,
)


class CompareAndSetStatus(str, Enum):
    """Outcome of a conditional intelligent-evaluation write."""

    APPLIED = "applied"
    NOT_FOUND = "not_found"
    CONFLICT = "conflict"


@dataclass(frozen=True)
class CompareAndSetResult:
    """Typed result for a status update guarded by an expected status."""

    status: CompareAndSetStatus
    evaluation: Optional[IntelligentEval] = None

    @property
    def applied(self) -> bool:
        return self.status is CompareAndSetStatus.APPLIED


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 list_page(self, offset: int, limit: int) -> list[IntelligentEval]:
        """Return one page of evals (created_at desc) with the same ordering as ``list_all``."""
        statement = (
            select(IntelligentEvalDB)
            .order_by(IntelligentEvalDB.created_at.desc())
            .offset(offset)
            .limit(limit)
        )
        return [self._from_db(r) for r in self.session.exec(statement).all()]

    def count(self) -> int:
        """Total number of evaluations (for pagination metadata)."""
        return self.session.exec(select(func.count()).select_from(IntelligentEvalDB)).one()

    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 get_with_sessions(
        self,
        eval_id: str,
    ) -> tuple[IntelligentEval, list[IntelligentEvalSession]] | None:
        """Load an evaluation and all session summaries in one DB snapshot."""

        statement = (
            select(IntelligentEvalDB, IntelligentEvalSessionDB)
            .join(
                IntelligentEvalSessionDB,
                IntelligentEvalSessionDB.eval_id == IntelligentEvalDB.id,
                isouter=True,
            )
            .where(IntelligentEvalDB.id == eval_id)
            .order_by(IntelligentEvalSessionDB.created_at.asc())
        )
        rows = self.session.exec(statement).all()
        if not rows:
            return None
        evaluation = self._from_db(rows[0][0])
        session_repo = IntelligentEvalSessionRepository(self.session)
        sessions = [session_repo._from_db(session_db) for _, session_db in rows if session_db is not None]
        return evaluation, sessions

    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 _compare_and_set_status(
        self,
        eval_id: str,
        *,
        expected_status: IntelligentEvalStatus,
        new_status: IntelligentEvalStatus,
    ) -> CompareAndSetResult:
        """Atomically transition status only when the expected state still holds."""
        now = utc_now()
        values = {
            "status": new_status.value,
            "updated_at": now,
        }
        if new_status is IntelligentEvalStatus.EXECUTING:
            values["started_at"] = func.coalesce(IntelligentEvalDB.started_at, now)
        if new_status in {
            IntelligentEvalStatus.COMPLETED,
            IntelligentEvalStatus.CANCELLED,
            IntelligentEvalStatus.FAILED,
        }:
            values["completed_at"] = now

        return self._compare_and_set_fields(
            eval_id,
            expected_status=expected_status,
            values=values,
        )

    def _compare_and_set_fields(
        self,
        eval_id: str,
        *,
        expected_status: IntelligentEvalStatus,
        values: dict[str, object],
    ) -> CompareAndSetResult:
        """Apply one conditional SQL update and classify its outcome."""
        try:
            result = self.session.exec(
                sql_update(IntelligentEvalDB)
                .where(
                    IntelligentEvalDB.id == eval_id,
                    IntelligentEvalDB.status == expected_status.value,
                )
                .values(**values)
            )
            if result.rowcount != 1:
                self.session.rollback()
                current = self.session.get(IntelligentEvalDB, eval_id)
                status = CompareAndSetStatus.NOT_FOUND if current is None else CompareAndSetStatus.CONFLICT
                return CompareAndSetResult(status=status)
            self.session.commit()
        except Exception:
            self.session.rollback()
            raise

        db = self.session.get(IntelligentEvalDB, eval_id)
        if db is None:
            return CompareAndSetResult(status=CompareAndSetStatus.NOT_FOUND)
        return CompareAndSetResult(status=CompareAndSetStatus.APPLIED, evaluation=self._from_db(db))

    def _submit_plan_if_planning(self, eval_id: str, plan: dict) -> CompareAndSetResult:
        now = utc_now()
        return self._compare_and_set_fields(
            eval_id,
            expected_status=IntelligentEvalStatus.PLANNING,
            values={
                "status": IntelligentEvalStatus.PENDING_APPROVAL.value,
                "plan": json.dumps(plan, ensure_ascii=False),
                "plan_feedback": None,
                "updated_at": now,
            },
        )

    def _reject_plan_if_pending(self, eval_id: str, feedback: str) -> CompareAndSetResult:
        now = utc_now()
        return self._compare_and_set_fields(
            eval_id,
            expected_status=IntelligentEvalStatus.PENDING_APPROVAL,
            values={
                "status": IntelligentEvalStatus.PLANNING.value,
                "plan_feedback": feedback,
                "updated_at": now,
            },
        )

    def _submit_report_if_executing(self, eval_id: str, report: dict) -> CompareAndSetResult:
        now = utc_now()
        return self._compare_and_set_fields(
            eval_id,
            expected_status=IntelligentEvalStatus.EXECUTING,
            values={
                "status": IntelligentEvalStatus.COMPLETED.value,
                "report": json.dumps(report, ensure_ascii=False),
                "updated_at": now,
                "completed_at": now,
            },
        )

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 list_by_evals(self, eval_ids: Sequence[str]) -> dict[str, list[IntelligentEvalSession]]:
        """Load sessions for several evaluations with one query.

        The result contains an empty list for every requested evaluation. This
        keeps read-model assembly deterministic while avoiding one count query
        per row in the evaluation list.
        """

        grouped = {eval_id: [] for eval_id in eval_ids}
        if not grouped:
            return grouped
        statement = (
            select(IntelligentEvalSessionDB)
            .where(IntelligentEvalSessionDB.eval_id.in_(tuple(grouped)))
            .order_by(IntelligentEvalSessionDB.created_at.asc())
        )
        for row in self.session.exec(statement).all():
            grouped.setdefault(row.eval_id, []).append(self._from_db(row))
        return grouped

    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 _create_if_executing(
        self,
        obj: IntelligentEvalSession,
    ) -> tuple[CompareAndSetStatus, Optional[IntelligentEvalSession]]:
        """Insert a session only while its parent evaluation is executing."""
        session_id = obj.id or new_uuid()
        now = obj.created_at or utc_now()
        statement = sql_insert(IntelligentEvalSessionDB).from_select(
            [
                "id",
                "eval_id",
                "target_id",
                "persona",
                "goal",
                "dimension",
                "status",
                "turn_count",
                "created_at",
            ],
            select(
                literal(session_id),
                IntelligentEvalDB.id,
                IntelligentEvalDB.target_id,
                literal(json.dumps(obj.persona, ensure_ascii=False)),
                literal(obj.goal),
                literal(obj.dimension),
                literal(IntelligentEvalSessionStatus.RUNNING.value),
                literal(0),
                literal(now),
            ).where(
                IntelligentEvalDB.id == obj.eval_id,
                IntelligentEvalDB.status == IntelligentEvalStatus.EXECUTING.value,
            ),
        )
        try:
            result = self.session.exec(statement)
            if result.rowcount != 1:
                self.session.rollback()
                evaluation = self.session.get(IntelligentEvalDB, obj.eval_id)
                status = CompareAndSetStatus.NOT_FOUND if evaluation is None else CompareAndSetStatus.CONFLICT
                return status, None
            self.session.commit()
        except Exception:
            self.session.rollback()
            raise

        db = self.session.get(IntelligentEvalSessionDB, session_id)
        if db is None:
            return CompareAndSetStatus.NOT_FOUND, None
        return CompareAndSetStatus.APPLIED, self._from_db(db)

    def _close_if_running(
        self,
        session_id: str,
        verdict: dict,
        status: IntelligentEvalSessionStatus = IntelligentEvalSessionStatus.COMPLETED,
    ) -> tuple[CompareAndSetStatus, Optional[IntelligentEvalSession]]:
        """Close exactly one running session without overwriting a race winner."""
        now = utc_now()
        try:
            result = self.session.exec(
                sql_update(IntelligentEvalSessionDB)
                .where(
                    IntelligentEvalSessionDB.id == session_id,
                    IntelligentEvalSessionDB.status == IntelligentEvalSessionStatus.RUNNING.value,
                )
                .values(
                    status=status.value,
                    verdict=json.dumps(verdict, ensure_ascii=False),
                    closed_at=now,
                )
            )
            if result.rowcount != 1:
                self.session.rollback()
                current = self.session.get(IntelligentEvalSessionDB, session_id)
                outcome = CompareAndSetStatus.NOT_FOUND if current is None else CompareAndSetStatus.CONFLICT
                return outcome, None
            self.session.commit()
        except Exception:
            self.session.rollback()
            raise

        db = self.session.get(IntelligentEvalSessionDB, session_id)
        if db is None:
            return CompareAndSetStatus.NOT_FOUND, None
        return CompareAndSetStatus.APPLIED, self._from_db(db)

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)

    def _create_user_and_increment(self, obj: IntelligentEvalMessage) -> CompareAndSetStatus:
        """Persist the sent user message and consume one turn atomically."""
        session_db = self.session.get(IntelligentEvalSessionDB, obj.session_id)
        if session_db is None:
            return CompareAndSetStatus.NOT_FOUND
        if session_db.status != IntelligentEvalSessionStatus.RUNNING.value:
            return CompareAndSetStatus.CONFLICT

        message_db = IntelligentEvalMessageDB(
            id=obj.id,
            session_id=obj.session_id,
            role="user",
            content=obj.content,
            latency_ms=None,
            created_at=obj.created_at,
        )
        try:
            session_db.turn_count += 1
            self.session.add(session_db)
            self.session.add(message_db)
            self.session.commit()
            self.session.refresh(message_db)
        except Exception:
            self.session.rollback()
            raise
        obj.id = message_db.id
        return CompareAndSetStatus.APPLIED