Spaces:
Sleeping
Sleeping
| 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 | |
| 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 | |
| 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") | |