models.py

"""
SQLEnv Pydantic models — the data contracts between client and server.

These models define the typed interface for the SQLEnv RL environment,
following the OpenEnv pattern (see OpenEnv Tutorial for reference):

    Action      — what the agent sends each step
    Observation — what the agent receives back
    State       — episode metadata (exposed via the state endpoint)

RL terminology — state vs observation
─────────────────────────────────────
In RL theory:

    State (s)       A COMPLETE description of the world. Nothing is hidden.
    Observation (o) A PARTIAL description of a state, which may omit info.

In SQLEnv these map to:

    EpisodeContext  The full RL state (s). Lives on the server only.
                    Contains gold answers, reward accumulators, DB
                    connection, full query history — everything needed
                    to advance the simulation and compute rewards.

    SQLObservation  The observation (o). Sent to the agent over the wire.
                    Contains the question, truncated results, revealed
                    schema, budget, and action history. The agent NEVER
                    sees the gold answer, progress scores, or full DB.

    SQLState        OpenEnv's "State" base class — lightweight episode
                    metadata (episode_id, step_count). This is NOT the
                    RL state; it is a convenience for logging/debugging.

This separation is what makes SQLEnv a POMDP: the agent must act under
uncertainty, which is what makes exploration necessary and learnable.
"""

import sqlite3
from dataclasses import dataclass, field as dataclass_field

from openenv.core.env_server.interfaces import Message
from openenv.core.env_server.types import Action, Observation, State
from pydantic import Field

# ---------------------------------------------------------------------------
# Wire types: these cross the HTTP boundary between client and server
# ---------------------------------------------------------------------------


class SQLAction(Action):
    """What the agent sends each step.

    The action space is intentionally small and structured so agents can
    explicitly control the environment loop.
    """

    action_type: str = Field(
        ...,
        description="One of: DESCRIBE, SAMPLE, QUERY, ANSWER",
    )
    argument: str = Field(
        ...,
        description=(
            "Table name (DESCRIBE/SAMPLE), SQL string (QUERY), "
            "or answer value (ANSWER)."
        ),
    )


class SQLObservation(Observation):
    """What the agent receives after each step.

    This is the agent's PARTIAL view of the world. Key design choices:

    - schema_info starts with table names only; columns are revealed
      incrementally as the agent DESCRIBEs tables.
    - result is always a truncated string, never raw data. The agent sees
      what a human analyst would see in a terminal — at most N rows of
      formatted text. This keeps the observation bounded and forces the
      agent to reason about what it sees rather than brute-force scanning.
    - action_history gives the agent memory of its own trajectory without
      the server needing to re-send full results from prior steps.
    """

    # Inherited from Observation: done (bool), reward (float | None)
    question: str = Field(..., description="The NL question to answer")
    schema_info: str = Field(..., description="Known schema information")
    result: str = Field(default="", description="Result of the last action")
    error: str = Field(default="", description="Error message if action failed")
    step_count: int = Field(default=0, description="Current step number")
    budget_remaining: int = Field(default=0, description="Steps remaining")
    action_history: list[str] = Field(
        default_factory=list,
        description="Summary of previous actions",
    )


class SQLState(State):
    """Episode metadata exposed via GET /state.

    This is the minimal public state — enough for logging and debugging,
    but NOT the full internal bookkeeping (see EpisodeContext below).
    """

    # # Inherited from State: episode_id (str | None), step_count (int)
    # game_name: str = Field(
    #     "sql_env", description="Name of the game/environment"
    # )
    history_messages: list[Message] = Field(default_factory=list)
    current_action_type: str = Field(
        default="QUERY",
        description="Current action type: DESCRIBE, SAMPLE, QUERY, or ANSWER",
    )


@dataclass
class QuestionRecord:
    """One question from the Spider dataset."""

    question_id: str
    question_text: str
    database_name: str
    gold_sql: str
    gold_answer: str
    answer_type: str
    difficulty: str
    tables_involved: list[str]


@dataclass
class EpisodeContext:
    """Per-episode server-side state (never sent to agent)."""

    episode_id: str
    db_connection: sqlite3.Connection
    question_record: QuestionRecord
    step_count: int = 0
    budget: int = 15
    described_tables: set[str] = dataclass_field(default_factory=set)
    action_log: list[str] = dataclass_field(default_factory=list)
    done: bool = False
    gold_answer: str | None = None
    gold_rows: list[tuple] = dataclass_field(default_factory=list)
    query_hashes: set[str] = dataclass_field(default_factory=set)
    previous_progress: float = 0.0