Delete server/environment.py

server/environment.py

@@ -1,619 +0,0 @@
-"""SQL Quality Environment - Core environment manager."""
-
-from __future__ import annotations
-
-import random as _random
-import sqlite3
-import threading
-import uuid
-from typing import Any, Dict
-
-from fastapi import HTTPException
-
-from .models import (
-    ActionType,
-    SQLAction,
-    SQLObservation,
-    SQLReward,
-    SQLState,
-    StepResponse,
-    TrajectoryStep,
-)
-from .tasks import (
-    create_dedup_bug_task,
-    create_fk_integrity_task,
-    create_null_audit_task,
-    create_schema_drift_task,
-    create_type_repair_task,
-)
-
-TASK_CREATORS = {
-    "null_audit": create_null_audit_task,
-    "type_repair": create_type_repair_task,
-    "dedup_bug": create_dedup_bug_task,
-    "fk_integrity": create_fk_integrity_task,
-    "schema_drift": create_schema_drift_task,
-}
-
-
-class SQLQualityEnvironment:
-    """Main environment class for SQL Quality tasks."""
-
-    def __init__(self) -> None:
-        """Initialize environment."""
-        self._current_task_id: str | None = None
-        self._episode_id: str | None = None
-        self._attempt_number: int = 1
-        self._step_count: int = (
-            0  # increments on every action (QUERY, PROFILE, FIX, SUBMIT)
-        )
-        self._max_steps_per_episode: int = 20  # prevent infinite loops
-        self._cumulative_reward: float = 0.0
-        self._done: bool = False
-        self._ground_truth: Dict[str, Any] = {}
-        self._cached_observation: SQLObservation | None = None
-        self._lock = threading.Lock()
-        # NEW: Multi-step trajectory support
-        self._trajectory: list = []
-        self._original_snapshot: dict = {}
-        self._original_conn: sqlite3.Connection | None = None
-        self._active: bool = False
-        self._last_feedback: str | None = None
-
-    def reset(self, task_id: str, seed: int | None = None) -> SQLObservation:
-        """Reset environment for a new episode.
-
-        Args:
-            task_id: One of the registered task IDs
-            seed: Optional deterministic episode seed
-
-        Returns:
-            SQLObservation for the new episode
-
-        Raises:
-            ValueError: If task_id is not valid
-        """
-        with self._lock:
-            # ── NEW: close any live connection from the previous episode ──────
-            if hasattr(self, "_original_conn") and self._original_conn is not None:
-                try:
-                    self._original_conn.close()
-                except Exception:
-                    pass  # Already closed or never opened — safe to ignore
-                self._original_conn = None
-
-            # Validate task_id
-            if task_id not in TASK_CREATORS:
-                raise ValueError(f"Invalid task_id: {task_id}")
-
-            # Generate new episode
-            self._episode_id = str(uuid.uuid4())
-            self._attempt_number = 1
-            self._step_count = (
-                0  # increments on every action (QUERY, PROFILE, FIX, SUBMIT)
-            )
-            self._cumulative_reward = 0.0
-            self._done = False
-            self._current_task_id = task_id
-            self._trajectory = []
-            self._original_snapshot = {}
-            self._ground_truth = {}
-            self._cached_observation = None
-            self._last_feedback = None
-
-            # Create task-specific database to extract observation data
-            conn = None
-            try:
-                # Generate random episode seed for diversity
-                episode_seed = seed if seed is not None else _random.randint(0, 2**31 - 1)
-                creator = TASK_CREATORS[task_id]
-                conn, observation, self._ground_truth = creator(episode_seed)
-
-                # Store original connection for graders (read-only use)
-                self._original_conn = conn
-
-                # Snapshot the target table for data preservation grading
-                target_table = self._ground_truth.get("table", "")
-                if target_table:
-                    from .graders import _snapshot_table
-
-                    self._original_snapshot = _snapshot_table(conn, target_table)
-
-                # Mark episode as active
-                self._active = True
-
-                # Cache observation for reuse in step()
-                self._cached_observation = observation.model_copy(
-                    update={"episode_id": self._episode_id}
-                )
-                
-                # Store schema info for dynamic updates
-                self._schema_info = observation.schema_info
-
-                # Keep connection open for multi-step operations
-                # Connection will be closed when episode ends or environment is reset
-
-            except Exception as e:
-                if conn is not None:
-                    conn.close()
-                raise e
-
-            return self._cached_observation
-
-    def step(self, action: SQLAction) -> StepResponse:
-        """
-        Execute one agent action and return the resulting observation and reward.
-
-        Routes by action_type:
-            QUERY   — execute SELECT on episode conn, return result in observation
-            PROFILE — compute column statistics, return in observation
-            FIX     — execute on fresh_test_db(), score partially, continue episode
-            SUBMIT  — call compute_episode_score(), end episode
-        """
-        with self._lock:
-            if not self._active:
-                raise HTTPException(
-                    status_code=422, detail="No active episode. Call reset() first."
-                )
-            if self._done:
-                raise HTTPException(
-                    status_code=422,
-                    detail="Episode is already done. Call reset() to start a new episode.",
-                )
-
-            # Check max steps guard
-            if self._step_count >= self._max_steps_per_episode:
-                raise HTTPException(
-                    status_code=422,
-                    detail="Max investigation steps reached. Call SUBMIT or reset().",
-                )
-
-            action = self._normalize_action(action)
-
-            action_type = action.action_type
-
-            # --- QUERY action ---
-            if action_type == ActionType.QUERY:
-                return self._handle_query(action)
-
-            # --- PROFILE action ---
-            elif action_type == ActionType.PROFILE:
-                return self._handle_profile(action)
-
-            # --- FIX action ---
-            elif action_type == ActionType.FIX:
-                return self._handle_fix(action)
-
-            # --- SUBMIT action (or legacy single-step) ---
-            else:
-                return self._handle_submit(action)
-
-    def _normalize_action(self, action: SQLAction) -> SQLAction:
-        """Infer v2 action_type for backward-compatible payloads without action_type."""
-        if "action_type" in action.model_fields_set:
-            return action
-
-        inferred = ActionType.SUBMIT
-        if action.column_name:
-            inferred = ActionType.PROFILE
-        elif action.fix_query:
-            inferred = ActionType.FIX
-        elif action.sql:
-            inferred = (
-                ActionType.QUERY
-                if action.sql.strip().upper().startswith("SELECT")
-                else ActionType.FIX
-            )
-        elif action.answer or action.root_cause_table or action.offending_key:
-            inferred = ActionType.SUBMIT
-
-        return action.model_copy(update={"action_type": inferred})
-
-    def _handle_query(self, action: SQLAction) -> StepResponse:
-        """
-        Handle a QUERY action.
-
-        Executes the SQL on the episode connection (supports multi-statement).
-        Returns a result preview in the observation. Does NOT end the episode.
-        Awards small step_reward (0.05) for queries that target the right area.
-        """
-        import re
-
-        sql = (action.sql or "").strip()
-
-        # Strip comments and leading whitespace to get the first keyword
-        sql_stripped = re.sub(r"--[^\n]*", "", sql)  # remove single-line comments
-        sql_stripped = re.sub(r"/\*.*?\*/", "", sql_stripped, flags=re.DOTALL)  # block comments
-        first_keyword = sql_stripped.strip().split()[0].upper() if sql_stripped.strip() else ""
-
-        # Allow all SQL types in query box for testing
-        # Only FIX action should be restricted
-        ALLOWED_READ_KEYWORDS = {"SELECT", "WITH", "EXPLAIN", "PRAGMA", "UPDATE", "INSERT", "DELETE", "CREATE", "DROP", "ALTER"}
-        if first_keyword not in ALLOWED_READ_KEYWORDS:
-            rejection_feedback = (
-                f"Invalid SQL keyword: '{first_keyword}'. "
-                f"Allowed keywords: {', '.join(sorted(ALLOWED_READ_KEYWORDS))}"
-            )
-            # Add to trajectory as a step but with zero reward and clear feedback
-            self._trajectory.append(
-                TrajectoryStep(
-                    step_number=len(self._trajectory) + 1,
-                    action_type=ActionType.QUERY,
-                    sql_submitted=sql[:200],
-                    result_preview=rejection_feedback,
-                    step_reward=0.0,
-                )
-            )
-            self._step_count = getattr(self, "_step_count", 0) + 1
-            obs = self._build_observation(result_preview=rejection_feedback)
-            return StepResponse(
-                observation=obs,
-                reward=SQLReward(
-                    score=self._cumulative_reward,
-                    step_reward=0.0,
-                    feedback=rejection_feedback,
-                ),
-                done=False,
-                info={"action_type": "query", "error": "non_select_rejected", "step": len(self._trajectory)},
-            )
-
-        # Split multiple statements by semicolon, but preserve semicolons in strings
-        statements = []
-        current = ""
-        in_string = False
-        string_char = None
-        i = 0
-        while i < len(sql):
-            char = sql[i]
-            if char in ("'", '"') and (i == 0 or sql[i-1] != "\\"):
-                if not in_string:
-                    in_string = True
-                    string_char = char
-                elif char == string_char:
-                    in_string = False
-                    string_char = None
-            if char == ";" and not in_string:
-                stmt = current.strip()
-                if stmt:
-                    statements.append(stmt)
-                current = ""
-            else:
-                current += char
-            i += 1
-        # Add final statement
-        if current.strip():
-            statements.append(current.strip())
-
-        # If no statements found, treat as single statement
-        if not statements:
-            statements = [sql] if sql else []
-
-        result_preview = ""
-        step_reward = 0.0
-        last_error = None
-
-        for stmt in statements:
-            if not stmt:
-                continue
-            try:
-                cursor = self._original_conn.execute(stmt)
-                # Check if query returns rows (SELECT) or just modifies (INSERT/UPDATE/DELETE)
-                if cursor.description:
-                    rows = cursor.fetchmany(5)
-                    cols = [d[0] for d in cursor.description]
-                    
-                    # Update schema info if this is a custom query with new columns
-                    if hasattr(self, '_schema_info') and self._schema_info:
-                        # Check if columns differ from current schema
-                        current_cols = []
-                        for table_info in self._schema_info:
-                            for col_str in table_info.columns:
-                                if isinstance(col_str, str):
-                                    current_cols.append(col_str.split(' (')[0])
-                        
-                        if set(cols) != set(current_cols):
-                            # Update schema with new columns from custom query
-                            from .models import SchemaInfo
-                            table_name = f"custom_query_{len(self._trajectory)}"
-                            new_columns = [f"{col} (TEXT)" for col in cols]
-                            self._schema_info.append(
-                                SchemaInfo(table_name=table_name, columns=new_columns, row_count=len(rows))
-                            )
-                    
-                    # Return actual data for formatting, not string
-                    result_preview = {
-                        "type": "table",
-                        "columns": cols,
-                        "rows": [dict(zip(cols, row)) for row in rows],
-                        "total_rows": len(rows)
-                    }
-                else:
-                    # For INSERT/UPDATE/DELETE, show affected rows
-                    result_preview = {
-                        "type": "affected_rows",
-                        "message": f"Query executed successfully. Rows affected: {cursor.rowcount}",
-                        "rows_affected": cursor.rowcount
-                    }
-            except sqlite3.Error as e:
-                last_error = str(e)
-                result_preview = f"SQL Error: {e}"
-                break  # Stop on first error
-
-        # Small reward for targeted queries (only check first statement for relevance)
-        if statements and not last_error:
-            from .graders import _query_reveals_issue
-            step_reward = (
-                0.05 if _query_reveals_issue(statements[0], self._ground_truth) else 0.01
-            )
-
-        # Add to trajectory
-        self._trajectory.append(
-            TrajectoryStep(
-                step_number=len(self._trajectory) + 1,
-                action_type=ActionType.QUERY,
-                sql_submitted=sql,
-                result_preview=str(result_preview)[:200] if isinstance(result_preview, (dict, list)) else str(result_preview)[:200],
-                step_reward=step_reward,
-            )
-        )
-
-        self._cumulative_reward = min(
-            1.0, round(self._cumulative_reward + step_reward, 4)
-        )
-        self._last_feedback = str(result_preview)[:200] if isinstance(result_preview, (dict, list)) else str(result_preview)[:200]
-        self._step_count += 1
-
-        # Build updated observation with trajectory history
-        obs = self._build_observation(result_preview=result_preview)
-
-        return StepResponse(
-            observation=obs,
-            reward=SQLReward(
-                score=self._cumulative_reward,
-                step_reward=step_reward,
-                feedback=str(result_preview)[:200] if isinstance(result_preview, (dict, list)) else str(result_preview)[:200],
-            ),
-            done=False,
-            info={"action_type": "query", "step": len(self._trajectory)},
-        )
-
-    def _handle_profile(self, action: SQLAction) -> StepResponse:
-        """
-        Handle a PROFILE action — statistical summary of a named column.
-
-        Computes: null rate, distinct count, data type distribution, min/max.
-        Does NOT end the episode. Awards step_reward based on column relevance.
-        """
-        col = (action.column_name or "").strip()
-        result_preview = ""
-
-        if not col:
-            result_preview = "Error: column_name required for PROFILE action."
-        else:
-            # Try to find which table has this column
-            try:
-                tables = self._original_conn.execute(
-                    "SELECT name FROM sqlite_master WHERE type='table'"
-                ).fetchall()
-                profile_results = []
-                for (tbl,) in tables:
-                    try:
-                        cols = [
-                            row[1]
-                            for row in self._original_conn.execute(
-                                f"PRAGMA table_info({tbl})"
-                            ).fetchall()
-                        ]
-                        if col in cols:
-                            total = self._original_conn.execute(
-                                f"SELECT COUNT(*) FROM {tbl}"
-                            ).fetchone()[0]
-                            nulls = self._original_conn.execute(
-                                f"SELECT COUNT(*) FROM {tbl} WHERE {col} IS NULL"
-                            ).fetchone()[0]
-                            distinct = self._original_conn.execute(
-                                f"SELECT COUNT(DISTINCT {col}) FROM {tbl}"
-                            ).fetchone()[0]
-                            profile_results.append(
-                                f"{tbl}.{col}: total={total}, nulls={nulls} ({100*nulls//total if total else 0}%), distinct={distinct}"
-                            )
-                    except sqlite3.Error:
-                        pass
-                result_preview = (
-                    " | ".join(profile_results)
-                    if profile_results
-                    else f"Column '{col}' not found."
-                )
-            except sqlite3.Error as e:
-                result_preview = f"Profile error: {e}"
-
-        # Determine step_reward based on relevance to ground truth
-        ground_table = self._ground_truth.get("table", "").lower()
-        ground_col = self._ground_truth.get(
-            "column", self._ground_truth.get("null_column", "")
-        ).lower()
-        col_lower = col.lower()
-
-        if col_lower == ground_col:
-            step_reward = 0.08  # On-target: profiling exact relevant column
-        elif ground_table and ground_table in result_preview.lower():
-            step_reward = 0.05  # Near-target: profiled something from right table
-        elif col_lower in result_preview.lower() and len(col_lower) > 2:
-            step_reward = 0.01  # Off-target: profiled something real but irrelevant
-        else:
-            step_reward = 0.01  # Minimal: invalid column or junk probe
-
-        self._trajectory.append(
-            TrajectoryStep(
-                step_number=len(self._trajectory) + 1,
-                action_type=ActionType.PROFILE,
-                sql_submitted=f"PROFILE({col})",
-                result_preview=result_preview[:200],
-                step_reward=step_reward,
-            )
-        )
-
-        self._cumulative_reward = min(
-            1.0, round(self._cumulative_reward + step_reward, 4)
-        )
-        self._last_feedback = result_preview[:200]
-        self._step_count += 1
-        obs = self._build_observation(result_preview=result_preview)
-
-        return StepResponse(
-            observation=obs,
-            reward=SQLReward(
-                score=self._cumulative_reward,
-                step_reward=step_reward,
-                feedback=result_preview[:200],
-            ),
-            done=False,
-            info={"action_type": "profile", "column_profiled": col},
-        )
-
-    def _handle_fix(self, action: SQLAction) -> StepResponse:
-        """
-        Handle a FIX action — execute write query on isolated copy for preview scoring.
-
-        Executes on a _fresh_test_db() copy. Does NOT mutate episode state.
-        Returns intermediate feedback on whether fix appears to work.
-        Does NOT end episode — agent can follow up with verification query or SUBMIT.
-        Awards intermediate reward based on quick fix assessment.
-        """
-        fix_sql = action.effective_sql or ""
-        step_reward = 0.0
-        feedback = ""
-
-        if fix_sql:
-            from .graders import grade_fix_correctness
-
-            quick_score, severity = grade_fix_correctness(
-                fix_sql,
-                self._original_conn,
-                self._original_snapshot,
-                self._ground_truth,
-            )
-            step_reward = quick_score * 0.3  # Partial credit for fix attempt
-            feedback = f"Fix preview: {severity} ({quick_score:.2f}). Submit to finalize score."
-        else:
-            feedback = "No SQL provided in fix action."
-
-        self._trajectory.append(
-            TrajectoryStep(
-                step_number=len(self._trajectory) + 1,
-                action_type=ActionType.FIX,
-                sql_submitted=fix_sql,
-                result_preview=feedback[:200],
-                step_reward=step_reward,
-            )
-        )
-
-        self._cumulative_reward = min(
-            1.0, round(self._cumulative_reward + step_reward, 4)
-        )
-        self._last_feedback = feedback
-        self._step_count += 1
-        obs = self._build_observation(result_preview=feedback)
-
-        return StepResponse(
-            observation=obs,
-            reward=SQLReward(
-                score=self._cumulative_reward,
-                step_reward=step_reward,
-                feedback=feedback,
-            ),
-            done=False,
-            info={
-                "action_type": "fix",
-                "severity_preview": severity if fix_sql else "N/A",
-            },
-        )
-
-    def _handle_submit(self, action: SQLAction) -> StepResponse:
-        """
-        Handle a SUBMIT action — compute final composite score and end episode.
-
-        Calls compute_episode_score() with the full trajectory.
-        Sets done=True. This is the only action that ends the episode.
-        """
-        from .graders import compute_episode_score
-
-        reward = compute_episode_score(
-            trajectory=self._trajectory,
-            ground_truth=self._ground_truth,
-            original_conn=self._original_conn,
-            original_snapshot=self._original_snapshot,
-            submit_action=action,
-        )
-
-        self._done = True
-        self._active = False
-        self._cumulative_reward = min(
-            1.0, round(self._cumulative_reward + reward.score, 4)
-        )
-        self._last_feedback = reward.feedback
-        obs = self._build_observation()
-
-        # Close connection when episode ends
-        if self._original_conn:
-            self._original_conn.close()
-            self._original_conn = None
-
-        return StepResponse(
-            observation=obs,
-            reward=reward,
-            done=True,
-            info={
-                "action_type": "submit",
-                "final_score": reward.score,
-                "dimensions": reward.dimensions,
-                "trajectory_length": len(self._trajectory),
-            },
-        )
-
-    def _build_observation(self, result_preview: str = "") -> SQLObservation:
-        """Build current observation including trajectory history."""
-        history = [
-            {
-                "step": s.step_number,
-                "type": s.action_type.value,
-                "sql": (s.sql_submitted or "")[:80],
-                "result": s.result_preview[:80],
-                "reward": s.step_reward,
-            }
-            for s in self._trajectory[-5:]  # Last 5 steps only (context window)
-        ]
-        return SQLObservation(
-            task_id=self._current_task_id,
-            task_description=self._cached_observation.task_description,
-            schema_info=self._cached_observation.schema_info,
-            sample_rows=self._cached_observation.sample_rows,
-            attempt_number=self._attempt_number,
-            episode_id=self._episode_id,
-            trajectory_history=history,
-            investigation_hint=result_preview[:200] if result_preview else None,
-            previous_feedback=(
-                self._last_feedback if hasattr(self, "_last_feedback") else None
-            ),
-            max_attempts=self._cached_observation.max_attempts,
-        )
-
-    def state(self) -> SQLState:
-        """Get current environment state.
-
-        Returns:
-            SQLState with current episode information
-
-        Raises:
-            ValueError: If no episode has been started
-        """
-        if self._episode_id is None:
-            raise ValueError("No episode started. Call reset() first.")
-
-        return SQLState(
-            task_id=self._current_task_id,
-            episode_id=self._episode_id,
-            attempt_number=self._attempt_number,
-            cumulative_reward=self._cumulative_reward,
-            done=self._done,
-        )