"""Repository for evaluation runs.""" from dataclasses import dataclass from enum import Enum from typing import Optional from sqlalchemy import insert as sql_insert from sqlalchemy import literal from sqlalchemy.exc import IntegrityError from sqlmodel import select from agenteval.models import CampaignStatus, EvalResult, EvalRun, RunStatus, RunTrigger from agenteval.storage.db import ( CampaignDB, EvalResultDB, EvalRunDB, TurnDB, new_uuid, utc_now, ) from agenteval.storage.repository.base import BaseRepository from agenteval.storage.repository.result import result_from_db class CampaignRunClaimStatus(str, Enum): """Outcome of a durable Campaign child-Run claim.""" CLAIMED = "claimed" EXISTING = "existing" NOT_FOUND = "not_found" CONFLICT = "conflict" @dataclass(frozen=True) class CampaignRunClaimResult: """Typed result for the idempotent child-Run claim seam.""" status: CampaignRunClaimStatus run: Optional[EvalRun] = None @property def accepted(self) -> bool: return self.status in (CampaignRunClaimStatus.CLAIMED, CampaignRunClaimStatus.EXISTING) @property def created(self) -> bool: return self.status is CampaignRunClaimStatus.CLAIMED class RunRepository(BaseRepository[EvalRun, EvalRunDB]): """Repository for evaluation runs.""" _table = EvalRunDB _order_by = "started_at" def claim_campaign_run( self, *, campaign_id: str, scenario_id: str, scenario_version: int, plan_index: int, occurrence_index: int, ) -> "CampaignRunClaimResult": """Claim one durable child Run occurrence while its Campaign is running. The conditional ``INSERT ... SELECT`` makes the Campaign status check part of the write. The composite unique index turns retries or a competing scheduler into an ``EXISTING`` result instead of a second pending Run. No external messages are sent by this operation. """ if plan_index < 0 or occurrence_index < 0: raise ValueError("plan_index and occurrence_index must be non-negative") run_id = new_uuid() now = utc_now() statement = sql_insert(EvalRunDB).from_select( [ "id", "target_id", "scenario_id", "scenario_version", "campaign_id", "campaign_plan_index", "campaign_occurrence_index", "status", "triggered_by", "started_at", ], select( literal(run_id), CampaignDB.target_id, literal(scenario_id), literal(scenario_version), CampaignDB.id, literal(plan_index), literal(occurrence_index), literal(RunStatus.PENDING.value), literal(RunTrigger.CAMPAIGN.value), literal(now), ).where( CampaignDB.id == campaign_id, CampaignDB.status == CampaignStatus.RUNNING.value, CampaignDB.target_id.is_not(None), ), ) try: result = self.session.exec(statement) if result.rowcount == 1: self.session.commit() db = self.session.get(EvalRunDB, run_id) return CampaignRunClaimResult(CampaignRunClaimStatus.CLAIMED, self._from_db(db)) self.session.rollback() except IntegrityError: self.session.rollback() except Exception: self.session.rollback() raise existing = self._get_campaign_identity(campaign_id, plan_index, occurrence_index) campaign = self.session.get(CampaignDB, campaign_id) if campaign is None: return CampaignRunClaimResult(CampaignRunClaimStatus.NOT_FOUND) if campaign.status != CampaignStatus.RUNNING.value: return CampaignRunClaimResult(CampaignRunClaimStatus.CONFLICT) if existing is not None: return CampaignRunClaimResult(CampaignRunClaimStatus.EXISTING, self._from_db(existing)) return CampaignRunClaimResult(CampaignRunClaimStatus.CONFLICT) def _get_campaign_identity(self, campaign_id: str, plan_index: int, occurrence_index: int) -> Optional[EvalRunDB]: statement = select(EvalRunDB).where( EvalRunDB.campaign_id == campaign_id, EvalRunDB.campaign_plan_index == plan_index, EvalRunDB.campaign_occurrence_index == occurrence_index, ) return self.session.exec(statement).first() def _to_db(self, run: EvalRun) -> EvalRunDB: db = EvalRunDB( id=run.id, target_id=run.target_id, scenario_id=run.scenario_id, scenario_version=run.scenario_version, campaign_id=run.campaign_id, campaign_plan_index=run.campaign_plan_index, campaign_occurrence_index=run.campaign_occurrence_index, status=run.status.value, triggered_by=run.triggered_by.value, started_at=run.started_at, completed_at=run.completed_at, ) if run.summary is not None: db.set_summary(run.summary.model_dump(mode="json")) return db def _from_db(self, db: EvalRunDB) -> EvalRun: return EvalRun( id=db.id, target_id=db.target_id, scenario_id=db.scenario_id, scenario_version=db.scenario_version or 1, campaign_id=db.campaign_id, campaign_plan_index=db.campaign_plan_index, campaign_occurrence_index=db.campaign_occurrence_index, status=db.status, triggered_by=db.triggered_by or "manual", started_at=db.started_at, completed_at=db.completed_at, summary=db.get_summary(), ) def list_by_campaign(self, campaign_id: str) -> list[EvalRun]: statement = select(EvalRunDB).where(EvalRunDB.campaign_id == campaign_id).order_by(EvalRunDB.started_at) return [self._from_db(r) for r in self.session.exec(statement).all()] def list_by_campaigns(self, campaign_ids: list[str]) -> dict[str, list[EvalRun]]: """Load child Runs for multiple Campaigns in one query.""" grouped = {campaign_id: [] for campaign_id in campaign_ids} if not campaign_ids: return grouped statement = ( select(EvalRunDB) .where(EvalRunDB.campaign_id.in_(campaign_ids)) # type: ignore[union-attr] .order_by(EvalRunDB.started_at) ) for row in self.session.exec(statement).all(): if row.campaign_id in grouped: grouped[row.campaign_id].append(self._from_db(row)) return grouped def mark_orphans_failed(self) -> int: """Fail process-orphaned Runs while preserving recoverable child claims. Every RUNNING Run may already have sent messages and must not replay. A PENDING Run is preserved only when it has a complete Campaign plan identity and its parent Campaign is still running; the Campaign loop can safely resume those claims before any external send occurred. """ statement = select(EvalRunDB).where(EvalRunDB.status.in_(["running", "pending"])) # type: ignore[attr-defined] candidates = self.session.exec(statement).all() # Batch load all campaigns for pending runs to avoid N+1 queries campaign_ids = {db.campaign_id for db in candidates if db.campaign_id and db.status == RunStatus.PENDING.value} campaigns = {} if campaign_ids: campaign_stmt = select(CampaignDB).where(CampaignDB.id.in_(campaign_ids)) # type: ignore[attr-defined] campaigns = {c.id: c for c in self.session.exec(campaign_stmt).all()} failed: list[EvalRunDB] = [] for db in candidates: if db.status == RunStatus.PENDING.value and self._is_recoverable_campaign_pending(db, campaigns): continue self._set_interrupted(db) failed.append(db) if failed: self.session.commit() return len(failed) def _is_recoverable_campaign_pending(self, db: EvalRunDB, campaigns: dict[str, CampaignDB]) -> bool: if db.campaign_id is None or db.campaign_plan_index is None or db.campaign_occurrence_index is None: return False campaign = campaigns.get(db.campaign_id) return campaign is not None and campaign.status == CampaignStatus.RUNNING.value def _set_interrupted(self, db: EvalRunDB) -> None: db.status = RunStatus.FAILED.value db.completed_at = db.completed_at or utc_now() summary = db.get_summary() or {} summary["error"] = {"code": "interrupted", "message": "服务重启导致评测中断"} db.set_summary(summary) self.session.add(db) def list_pending_campaign_children(self, campaign_id: str) -> list[EvalRun]: """Return durable pending claims in deterministic plan order.""" statement = ( select(EvalRunDB) .where( EvalRunDB.campaign_id == campaign_id, EvalRunDB.campaign_plan_index.is_not(None), EvalRunDB.campaign_occurrence_index.is_not(None), EvalRunDB.status == RunStatus.PENDING.value, ) .order_by(EvalRunDB.campaign_plan_index, EvalRunDB.campaign_occurrence_index) ) return [self._from_db(db) for db in self.session.exec(statement).all()] def mark_campaign_running_interrupted(self, campaign_id: str) -> list[str]: """Fail running child Runs from a previous process without replaying.""" statement = select(EvalRunDB).where( EvalRunDB.campaign_id == campaign_id, EvalRunDB.campaign_plan_index.is_not(None), EvalRunDB.campaign_occurrence_index.is_not(None), EvalRunDB.status == RunStatus.RUNNING.value, ) rows = list(self.session.exec(statement).all()) for db in rows: self._set_interrupted(db) if rows: self.session.commit() return [db.id or "" for db in rows] def mark_pending_campaign_child_interrupted(self, run_id: str) -> Optional[EvalRun]: """Fail one still-pending child claim that cannot be safely recovered.""" db = self.session.get(EvalRunDB, run_id) if db is None or db.status != RunStatus.PENDING.value: return self._from_db(db) if db is not None else None self._set_interrupted(db) self.session.commit() self.session.refresh(db) return self._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.scenario_version = run.scenario_version existing.campaign_id = run.campaign_id existing.campaign_plan_index = run.campaign_plan_index existing.campaign_occurrence_index = run.campaign_occurrence_index existing.status = run.status.value existing.triggered_by = run.triggered_by.value existing.completed_at = run.completed_at if run.summary is not None: existing.set_summary(run.summary.model_dump(mode="json")) self.session.add(existing) self.session.commit() self.session.refresh(existing) return self._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()]