RonyForAI/Mirage_DB_RL

1from openenv.core.env_server.types import Action, Observation, State2from pydantic import Field3from typing import List4 5# ──────────────── ACTION ─────────────────────────────────────────────────────6class QueryAction(Action):7    next_table: int = Field(..., ge=0,   description="Index of next table to join (from remaining_tables)")8    join_type:  int = Field(..., ge=0, le=2, description="Join algorithm: 0=hash(1.0×) 1=nested-loop(2.0×) 2=merge-sort(0.8×)")9    use_index:  int = Field(..., ge=0, le=1, description="Index scan: 0=full-table-scan 1=index-scan(halves base rows if index present)")10 11# ──────────────── OBSERVATION ────────────────────────────────────────────────12class QueryObservation(Observation):13    # Schema visible to the agent14    tables:           List[str]    # table names in this query scenario15    table_rows:       List[int]    # cardinality estimates (noisy — simulates planner statistics)16    selectivities:    List[float]  # join predicate selectivities17    has_index:        List[int]    # index availability per table (1=yes, 0=no)18 19    # Episode progress20    chosen_order:     List[int]    # table indices already joined, in insertion order21    remaining_tables: List[int]    # table indices not yet joined (valid choices for next_table)22    step_number:      int          # 0-based step counter23    current_cost:     float        # accumulated join cost so far (based on true cardinalities)24    intermediate_size: float = Field(25        default=1.0,26        description=(27            "Estimated size of the intermediate result accumulated so far. "28            "Computed as the product of (est_rows × selectivity) for each joined table. "29            "Joining large-output tables early causes this to explode, compounding all "30            "subsequent join costs. Keep this small by joining selective tables first."31        ),32    )33 34    # Enterprise context35    query_context:    str = Field(default="", description="SQL-like description of the query being optimized")36 37# ──────────────── STATE ──────────────────────────────────────────────────────38class QueryState(State):39    chosen_order:     List[int]  = Field(default_factory=list)40    remaining_tables: List[int]  = Field(default_factory=list)41    current_cost:     float      = 0.042    final_cost:       float      = 0.043    scenario_name:    str        = ""   # which scenario was sampled this episode