Upload 2 files

server/main.py

@@ -65,12 +65,25 @@ async def root():
 
 @app.get("/health", response_model=HealthResponse)
 async def health() -> HealthResponse:
-    """Health check endpoint for HF Space ping validation.
+    """Health check endpoint for monitoring and load balancers.
 
     Returns:
         dict: Status response indicating the server is running
     """
-    return HealthResponse(status="ok")
+    # Check LLM availability
+    llm_available = False
+    try:
+        import os
+        hf_token = os.getenv("HF_TOKEN", "")
+        api_base_url = os.getenv("API_BASE_URL", "")
+        model_name = os.getenv("MODEL_NAME", "")
+        
+        if hf_token and api_base_url and model_name:
+            llm_available = True
+    except Exception:
+        pass
+    
+    return HealthResponse(status="ok", llm_available=llm_available)
 
 
 @app.post("/ping", response_model=HealthResponse)

server/models.py

@@ -0,0 +1,192 @@
+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.0, 1.0]."""
+        return round(max(0.0, min(1.0, 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 = Field(
+        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
+    info: 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")