"""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, ) @staticmethod def visible(): """可见性谓词:逻辑删除的评估默认不可见(查询位点共用接缝)。""" return IntelligentEvalDB.status != IntelligentEvalStatus.DELETED.value @staticmethod def _is_visible(db: IntelligentEvalDB) -> bool: return db.status != IntelligentEvalStatus.DELETED.value def list_all(self) -> list[IntelligentEval]: statement = select(IntelligentEvalDB).where(self.visible()).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, status: Optional[str] = None) -> list[IntelligentEval]: """Return one page of evals (created_at desc), optionally filtered by status. Deleted evals are always excluded — the list API has no recycle bin. """ statement = select(IntelligentEvalDB).where(self.visible()).order_by(IntelligentEvalDB.created_at.desc()) if status: statement = statement.where(IntelligentEvalDB.status == status) statement = statement.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).where(self.visible())).one() def count_by_status(self) -> dict[str, int]: """Count evaluations per status (for the list page stat bar).""" rows = self.session.exec( select(IntelligentEvalDB.status, func.count(IntelligentEvalDB.id)) .where(self.visible()) .group_by(IntelligentEvalDB.status) ).all() stats = {status.value: 0 for status in IntelligentEvalStatus if status != IntelligentEvalStatus.DELETED} for status, cnt in rows: stats[status] = cnt return stats def get(self, eval_id: str) -> Optional[IntelligentEval]: db = self.session.get(IntelligentEvalDB, eval_id) if db is None or not self._is_visible(db): return None return self._from_db(db) def require_live_eval(self, eval_id: str) -> IntelligentEval: """可见性单一接缝:存在且未删 → 返回实体,否则 ``LookupError``。""" ev = self.get(eval_id) if ev is None: raise LookupError(f"intelligent eval {eval_id} not found") return ev def get_including_deleted(self, eval_id: str) -> Optional[IntelligentEval]: """包含已删除评估的原样读取(供 delete 幂等检查使用)。""" 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, self.visible(), ) .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