Blame for server/models.py

from __future__ import annotations

from enum import Enum
from typing import Optional

from pydantic import BaseModel, Field, field_validator


class ActionType(str, Enum):
    """
    The four action types available to an agent in sql_quality_env.

    QUERY   — Read-only SELECT statement. Inspect schema or data.
              Does not modify the database. Can be called multiple times.
    PROFILE — Run null-rate and type distribution stats on a named column.
              Shortcut for common inspection patterns. Read-only.
    FIX     — Submit a write query intended to resolve the data issue.
              Executes on an isolated _fresh_test_db() copy for scoring.
    SUBMIT  — Declare final diagnosis and close the episode.
              Triggers composite score computation. Episode ends after this.
    """

    QUERY = "query"
    PROFILE = "profile"
    FIX = "fix"
    SUBMIT = "submit"


class SchemaInfo(BaseModel):
    """Schema information for a database table."""

    table_name: str
    columns: list[str]  # Format: "column_name (TYPE)"
    row_count: int


class SQLObservation(BaseModel):
    """What the agent sees at each step."""

    task_id: str = Field(
        description=(
            "One of: null_audit, type_repair, dedup_bug, fk_integrity, schema_drift"
        )
    )
    task_description: str = Field(
        description="Natural language description of the task"
    )
    schema_info: list[SchemaInfo] = Field(description="Schema of all relevant tables")
    sample_rows: dict[str, list[dict]] = Field(
        description="First 5 rows of each table as list of dicts"
    )
    episode_id: Optional[str] = Field(
        default=None, description="UUID4 identifying this episode"
    )
    attempt_number: int = Field(ge=1, description="Current attempt (1-indexed)")
    max_attempts: int = Field(
        ge=1, description="Maximum attempts allowed for this task"
    )
    previous_feedback: Optional[str] = Field(
        default=None, description="Feedback from previous attempt, if any"
    )
    # NEW FIELDS:
    trajectory_history: list[dict] = Field(
        default_factory=list,
        description=(
            "List of previous actions and their results in this episode. "
            "Each entry: {action_type, sql_or_answer, result_preview, step_reward}. "
            "Empty on first step after reset()."
        ),
    )
    available_actions: list[str] = Field(
        default_factory=lambda: ["query", "profile", "fix", "submit"],
        description="Action types available at this point in the episode.",
    )
    investigation_hint: str | None = Field(
        default=None,
        description=(
            "Hint provided after a failed attempt. None on first step. "
            "Populated with actionable feedback after each scored action."
        ),
    )


class SQLAction(BaseModel):
    """
    Agent action for sql_quality_env.

    Fields:
        task_id:       Which task this action targets.
        action_type:   One of QUERY / PROFILE / FIX / SUBMIT.
        sql:           For QUERY and FIX actions — the SQL statement to execute.
        column_name:   For PROFILE actions — the column to profile.
        answer:        For null_audit SUBMIT — the column name identified.
        fix_query:     Alias for sql in FIX context (kept for backward compat).
        root_cause_table:  For dedup_bug SUBMIT — the table with duplicates.
        offending_key:     For dedup_bug SUBMIT — the key column with duplicates.
    """

    task_id: str
    action_type: ActionType = ActionType.SUBMIT
    sql: str | None = None
    column_name: str | None = None
    answer: str | None = None
    fix_query: str | None = None
    root_cause_table: str | None = None
    offending_key: str | None = None

    @property
    def effective_sql(self) -> str | None:
        """Return sql or fix_query, whichever is set. fix_query takes precedence."""
        return self.fix_query or self.sql


class SQLReward(BaseModel):
    """
    Reward returned after each step.

    Fields:
        score:          Episode-level composite score 0.0–1.0.
        step_reward:    Immediate reward for this specific action (not cumulative).
                        Used for dense reward shaping in RL training.
        feedback:       Human-readable explanation of what was scored and why.
        dimensions:     Per-dimension breakdown. Present only after SUBMIT.
                        None for intermediate QUERY/PROFILE/FIX actions.
        partial_scores: Legacy field kept for backward compatibility.
    """

    score: float
    step_reward: float = 0.0
    feedback: str = ""
    dimensions: dict | None = None
    partial_scores: dict | None = None

    @field_validator("score", "step_reward")
    @classmethod
    def clamp_to_unit(cls, v: float) -> float:
        """Ensure scores stay in [0.001, 0.999] and round to 4 decimal places."""
        return round(max(0.001, min(0.999, v)), 4)


class TrajectoryStep(BaseModel):
    """A single step in the agent's investigation trajectory."""

    step_number: int
    action_type: ActionType
    sql_submitted: str | None = None
    answer_submitted: str | None = None
    result_preview: str = ""
    step_reward: float = 0.0
    timestamp_ms: int = 0


class SQLState(BaseModel):
    """Current episode state."""

    task_id: str
    episode_id: str = Field(description="UUID4 identifying this episode")
    attempt_number: int
    cumulative_reward: float = Field(ge=0.0)
    done: bool


class ResetRequest(BaseModel):
    """Request body for /reset endpoint."""

    task_id: str | None = Field(
        default=None,
        description=(
            "One of: null_audit, type_repair, dedup_bug, fk_integrity, schema_drift"
        )
    )
    seed: int | None = Field(
        default=None,
        description=(
            "Optional deterministic episode seed. If omitted, environment chooses a random seed."
        ),
    )


class StepResponse(BaseModel):
    """Response body for /step endpoint."""

    observation: SQLObservation
    reward: SQLReward
    done: bool = Field(default=False, description="Whether the episode is complete")
    info: dict = Field(default_factory=dict)


class HealthResponse(BaseModel):
    """Response body for health/ping endpoints."""

    status: str = Field(description="Service health status")
    llm_available: Optional[bool] = Field(default=None, description="LLM connection status")