AgentEvalTool/backend/agenteval/storage/exploration_repository.py

"""Repositories for virtual-user exploration sessions (探索式评测)."""

from typing import Optional

from sqlmodel import Session, select

from agenteval.exploration.models import ExplorationMessage, ExplorationSession, ExplorationSessionStatus
from agenteval.storage.db import ExplorationMessageDB, ExplorationSessionDB, get_session, utc_now
from agenteval.storage.repository import BaseRepository


class ExplorationSessionRepository(BaseRepository[ExplorationSession, ExplorationSessionDB]):
    """Repository for virtual-user exploration sessions (探索会话)."""

    _table = ExplorationSessionDB
    _order_by = "created_at"

    def _copy_mutable(self, db: ExplorationSessionDB, session_obj: ExplorationSession) -> None:
        db.goal = session_obj.goal
        db.status = session_obj.status.value
        db.triggered_by = session_obj.triggered_by.value
        db.turn_count = session_obj.turn_count
        db.error = session_obj.error
        db.closed_at = session_obj.closed_at
        db.set_persona(session_obj.persona)
        if session_obj.seed_ref is not None:
            db.set_seed_ref(session_obj.seed_ref)
        if session_obj.experience is not None:
            db.set_experience(session_obj.experience)
        if session_obj.judge_review is not None:
            db.set_judge_review(session_obj.judge_review)

    def _to_db(self, session_obj: ExplorationSession) -> ExplorationSessionDB:
        db = ExplorationSessionDB(
            id=session_obj.id,
            campaign_id=session_obj.campaign_id,
            target_id=session_obj.target_id,
            created_at=session_obj.created_at,
        )
        self._copy_mutable(db, session_obj)
        return db

    def _from_db(self, db: ExplorationSessionDB) -> ExplorationSession:
        return ExplorationSession(
            id=db.id,
            campaign_id=db.campaign_id,
            target_id=db.target_id,
            persona=db.get_persona(),
            goal=db.goal,
            seed_ref=db.get_seed_ref(),
            status=db.status,
            triggered_by=db.triggered_by,
            experience=db.get_experience(),
            judge_review=db.get_judge_review(),
            turn_count=db.turn_count,
            error=db.error,
            created_at=db.created_at,
            closed_at=db.closed_at,
        )

    def update(self, session_obj: ExplorationSession) -> Optional[ExplorationSession]:
        existing = self.session.get(ExplorationSessionDB, session_obj.id)
        if not existing:
            return None
        self._copy_mutable(existing, session_obj)
        self.session.add(existing)
        self.session.commit()
        self.session.refresh(existing)
        return self._from_db(existing)

    def list_by_campaign(self, campaign_id: str) -> list[ExplorationSession]:
        statement = (
            select(ExplorationSessionDB)
            .where(ExplorationSessionDB.campaign_id == campaign_id)
            .order_by(ExplorationSessionDB.created_at)
        )
        return [self._from_db(r) for r in self.session.exec(statement).all()]

    def expire_running_sessions(self, campaign_id: str) -> int:
        """Expire every still-running exploration session of a finalized campaign.

        Returns the number of sessions expired. Completed/failed sessions keep
        their evidence untouched.
        """
        expired = 0
        for session_obj in self.list_by_campaign(campaign_id):
            if session_obj.status != ExplorationSessionStatus.RUNNING:
                continue
            session_obj.status = ExplorationSessionStatus.EXPIRED
            session_obj.closed_at = utc_now()
            self.update(session_obj)
            expired += 1
        return expired


class ExplorationMessageRepository:
    """Append-only repository for exploration session chat rows."""

    def __init__(self, session: Optional[Session] = None):
        self.session = session or get_session()

    def save_message(self, message: ExplorationMessage) -> ExplorationMessage:
        db = ExplorationMessageDB(
            id=message.id,
            session_id=message.session_id,
            round_index=message.round_index,
            role=message.role,
            content=message.content,
            latency_ms=message.latency_ms,
            created_at=message.created_at,
        )
        self.session.add(db)
        self.session.commit()
        self.session.refresh(db)
        message.id = db.id
        return message

    def list_by_session(self, session_id: str) -> list[ExplorationMessage]:
        statement = (
            select(ExplorationMessageDB)
            .where(ExplorationMessageDB.session_id == session_id)
            .order_by(ExplorationMessageDB.created_at)
        )
        return [
            ExplorationMessage(
                id=r.id,
                session_id=r.session_id,
                round_index=r.round_index,
                role=r.role,
                content=r.content,
                latency_ms=r.latency_ms,
                created_at=r.created_at,
            )
            for r in self.session.exec(statement).all()
        ]