AgentEvalTool/backend/agenteval/storage/db/campaign_tables.py

"""Campaign and exploration tables."""

from datetime import datetime
from typing import Any, Optional

from sqlmodel import Field, SQLModel

from agenteval.storage.db.core import _json_dumps, _json_loads, new_uuid, utc_now


class CampaignDB(SQLModel, table=True):
    """Database table for evaluation campaigns (评估活动)."""

    __tablename__ = "campaigns"

    id: Optional[str] = Field(default_factory=new_uuid, primary_key=True)
    name: str
    target_id: Optional[str] = Field(default=None, foreign_key="eval_targets.id")
    window_seconds: int
    time_scale: float = 1.0
    plan: str = "[]"
    status: str = "planned"
    started_at: Optional[datetime] = None
    completed_at: Optional[datetime] = None
    created_at: Optional[datetime] = Field(default_factory=utc_now)
    summary: Optional[str] = None
    analysis_model_config_id: Optional[str] = None
    exploration_seeds: Optional[str] = None
    exploration_budget: Optional[str] = None
    last_patrolled_at: Optional[datetime] = None

    def get_plan(self) -> list[dict[str, Any]]:
        return _json_loads(self.plan)

    def set_plan(self, plan: list[dict[str, Any]]) -> None:
        self.plan = _json_dumps(plan)

    def get_summary(self) -> Optional[dict[str, Any]]:
        return _json_loads(self.summary) if self.summary else None

    def set_summary(self, summary: dict[str, Any]) -> None:
        self.summary = _json_dumps(summary)

    def get_exploration_seeds(self) -> Optional[dict[str, Any]]:
        return _json_loads(self.exploration_seeds) if self.exploration_seeds else None

    def set_exploration_seeds(self, seeds: dict[str, Any]) -> None:
        self.exploration_seeds = _json_dumps(seeds)

    def get_exploration_budget(self) -> Optional[dict[str, Any]]:
        return _json_loads(self.exploration_budget) if self.exploration_budget else None

    def set_exploration_budget(self, budget: dict[str, Any]) -> None:
        self.exploration_budget = _json_dumps(budget)


class CampaignAnalysisDB(SQLModel, table=True):
    """One row per campaign holding its intelligent analysis (智能分析) state."""

    __tablename__ = "campaign_analyses"

    id: Optional[str] = Field(default_factory=new_uuid, primary_key=True)
    campaign_id: str = Field(unique=True, index=True)
    status: str = "generating"
    result: Optional[str] = None
    model_config_id: Optional[str] = None
    error: Optional[str] = None
    triggered_by: str = "manual"
    recovery_attempts: int = 0
    created_at: Optional[datetime] = Field(default_factory=utc_now)
    updated_at: Optional[datetime] = Field(default_factory=utc_now)

    def get_result(self) -> Optional[dict[str, Any]]:
        return _json_loads(self.result) if self.result else None

    def set_result(self, result: dict[str, Any]) -> None:
        self.result = _json_dumps(result)


class CampaignPeriodComparisonDB(SQLModel, table=True):
    """One row per campaign holding its period-comparison (周期对比) state.

    Mirrors ``campaign_analyses`` (one row, upserted per campaign), plus the
    baseline campaign this comparison was computed against.
    """

    __tablename__ = "campaign_period_comparisons"

    id: Optional[str] = Field(default_factory=new_uuid, primary_key=True)
    campaign_id: str = Field(unique=True, index=True)
    baseline_campaign_id: str = Field(index=True)
    status: str = "generating"
    result: Optional[str] = None
    model_config_id: Optional[str] = None
    error: Optional[str] = None
    triggered_by: str = "manual"
    recovery_attempts: int = 0
    created_at: Optional[datetime] = Field(default_factory=utc_now)
    updated_at: Optional[datetime] = Field(default_factory=utc_now)

    def get_result(self) -> Optional[dict[str, Any]]:
        return _json_loads(self.result) if self.result else None

    def set_result(self, result: dict[str, Any]) -> None:
        self.result = _json_dumps(result)


class ExplorationSessionDB(SQLModel, table=True):
    """One virtual-user exploration session (探索会话) within a campaign.

    Independent entity, deliberately NOT an EvalRun: exploration outcomes feed
    the 体验判定 evidence line and must not pollute pass-rate semantics
    (ADR-0001 comparability / ADR-0002).
    """

    __tablename__ = "exploration_sessions"

    id: Optional[str] = Field(default_factory=new_uuid, primary_key=True)
    campaign_id: str = Field(index=True, foreign_key="campaigns.id")
    target_id: str = Field(foreign_key="eval_targets.id")
    persona: str = "{}"
    goal: str = ""
    seed_ref: Optional[str] = None
    status: str = "running"
    triggered_by: str = "auto"
    experience: Optional[str] = None
    judge_review: Optional[str] = None
    turn_count: int = 0
    error: Optional[str] = None
    created_at: Optional[datetime] = Field(default_factory=utc_now)
    closed_at: Optional[datetime] = None

    def get_persona(self) -> dict[str, Any]:
        return _json_loads(self.persona)

    def set_persona(self, persona: dict[str, Any]) -> None:
        self.persona = _json_dumps(persona)

    def get_seed_ref(self) -> Optional[dict[str, Any]]:
        return _json_loads(self.seed_ref) if self.seed_ref else None

    def set_seed_ref(self, seed_ref: dict[str, Any]) -> None:
        self.seed_ref = _json_dumps(seed_ref)

    def get_experience(self) -> Optional[dict[str, Any]]:
        return _json_loads(self.experience) if self.experience else None

    def set_experience(self, experience: dict[str, Any]) -> None:
        self.experience = _json_dumps(experience)

    def get_judge_review(self) -> Optional[dict[str, Any]]:
        return _json_loads(self.judge_review) if self.judge_review else None

    def set_judge_review(self, judge_review: dict[str, Any]) -> None:
        self.judge_review = _json_dumps(judge_review)


class ExplorationMessageDB(SQLModel, table=True):
    """One chat message inside an exploration session (role/content form).

    Parallel to ``eval_runs`` turns but never shared with them. A completed
    question-answer pair contributes two rows (user + assistant); the reply
    row carries ``latency_ms``.
    """

    __tablename__ = "exploration_messages"

    id: Optional[str] = Field(default_factory=new_uuid, primary_key=True)
    session_id: str = Field(index=True, foreign_key="exploration_sessions.id")
    round_index: int = 0
    role: str = "user"
    content: str = ""
    latency_ms: Optional[int] = None
    created_at: Optional[datetime] = Field(default_factory=utc_now)