Upload inference.py

server/inference.py

@@ -0,0 +1,356 @@
+"""Inference module for SQL Quality Environment.
+
+This module provides an LLM-based AI agent interface for solving SQL data quality tasks.
+It uses OpenAI API to generate SQL queries and fixes dynamically.
+"""
+
+import json
+import os
+import sys
+from typing import Any, Dict, List, Optional
+
+import requests
+from openai import OpenAI
+
+# Configuration from environment
+MAX_STEPS = int(os.getenv("MAX_STEPS", "10"))
+API_BASE_URL = os.getenv("API_BASE_URL", "https://api.openai.com/v1")
+MODEL_NAME = os.getenv("MODEL_NAME", "gpt-4o-mini")
+HF_TOKEN = os.getenv("HF_TOKEN", "")
+
+
+class SQLQualityAgent:
+    """LLM-based agent for solving SQL data quality tasks."""
+
+    def __init__(self, base_url: str = "http://localhost:7860"):
+        self.base_url = base_url.rstrip("/")
+        self.task_id: Optional[str] = None
+        self.episode_id: Optional[str] = None
+        self.step_count = 0
+        self.trajectory: List[Dict[str, Any]] = []
+        
+        # Initialize OpenAI client
+        api_key = HF_TOKEN or os.getenv("OPENAI_API_KEY", "")
+        self.client = OpenAI(
+            api_key=api_key,
+            base_url=API_BASE_URL if "openai" not in API_BASE_URL else None
+        )
+
+    def health_check(self) -> bool:
+        """Check if the environment is running."""
+        try:
+            r = requests.get(f"{self.base_url}/health", timeout=5)
+            return r.status_code == 200 and r.json().get("status") == "ok"
+        except Exception:
+            return False
+
+    def reset(self, task_id: str, seed: Optional[int] = None) -> Dict[str, Any]:
+        """Initialize a new episode for the given task."""
+        self.task_id = task_id
+        payload = {"task_id": task_id}
+        if seed is not None:
+            payload["seed"] = seed
+
+        r = requests.post(f"{self.base_url}/reset", json=payload, timeout=30)
+        r.raise_for_status()
+        data = r.json()
+        self.episode_id = data.get("episode_id")
+        return data
+
+    def query(self, sql: str) -> Dict[str, Any]:
+        """Execute a SELECT query to investigate the data."""
+        if not self.task_id:
+            raise ValueError("No active task. Call reset() first.")
+        
+        self.step_count += 1
+        if self.step_count > MAX_STEPS:
+            raise RuntimeError(f"Exceeded MAX_STEPS ({MAX_STEPS})")
+
+        r = requests.post(
+            f"{self.base_url}/step",
+            json={
+                "task_id": self.task_id,
+                "action_type": "query",
+                "sql": sql,
+            },
+            timeout=30,
+        )
+        r.raise_for_status()
+        return r.json()
+
+    def profile(self, column: str) -> Dict[str, Any]:
+        """Profile a column to get statistics."""
+        if not self.task_id:
+            raise ValueError("No active task. Call reset() first.")
+        
+        self.step_count += 1
+        if self.step_count > MAX_STEPS:
+            raise RuntimeError(f"Exceeded MAX_STEPS ({MAX_STEPS})")
+
+        r = requests.post(
+            f"{self.base_url}/step",
+            json={
+                "task_id": self.task_id,
+                "action_type": "profile",
+                "column_name": column,
+            },
+            timeout=30,
+        )
+        r.raise_for_status()
+        return r.json()
+
+    def fix(self, fix_sql: str) -> Dict[str, Any]:
+        """Apply a fix SQL (UPDATE/INSERT/DELETE)."""
+        if not self.task_id:
+            raise ValueError("No active task. Call reset() first.")
+        
+        self.step_count += 1
+        if self.step_count > MAX_STEPS:
+            raise RuntimeError(f"Exceeded MAX_STEPS ({MAX_STEPS})")
+
+        r = requests.post(
+            f"{self.base_url}/step",
+            json={
+                "task_id": self.task_id,
+                "action_type": "fix",
+                "fix_query": fix_sql,
+            },
+            timeout=30,
+        )
+        r.raise_for_status()
+        return r.json()
+
+    def submit(
+        self,
+        answer: Optional[str] = None,
+        fix_query: Optional[str] = None,
+        root_cause_table: Optional[str] = None,
+        offending_key: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """Submit final answer to complete the episode."""
+        if not self.task_id:
+            raise ValueError("No active task. Call reset() first.")
+
+        payload: Dict[str, Any] = {
+            "task_id": self.task_id,
+            "action_type": "submit",
+        }
+        if answer:
+            payload["answer"] = answer
+        if fix_query:
+            payload["fix_query"] = fix_query
+        if root_cause_table:
+            payload["root_cause_table"] = root_cause_table
+        if offending_key:
+            payload["offending_key"] = offending_key
+
+        r = requests.post(f"{self.base_url}/step", json=payload, timeout=30)
+        r.raise_for_status()
+        return r.json()
+
+
+    def llm_generate_sql(self, prompt: str, system_context: str = "") -> str:
+        """Generate SQL using LLM."""
+        messages = [
+            {"role": "system", "content": system_context or "You are a SQL data quality expert. Generate valid SQLite SQL queries."},
+            {"role": "user", "content": prompt}
+        ]
+        
+        try:
+            response = self.client.chat.completions.create(
+                model=MODEL_NAME,
+                messages=messages,
+                temperature=0.1,
+                max_tokens=500
+            )
+            sql = response.choices[0].message.content.strip()
+            # Extract SQL from markdown if present
+            if "```sql" in sql:
+                sql = sql.split("```sql")[1].split("```")[0].strip()
+            elif "```" in sql:
+                sql = sql.split("```")[1].split("```")[0].strip()
+            return sql
+        except Exception as e:
+            print(f"LLM generation failed: {e}")
+            return ""
+
+
+def run_agent(task_id: str, base_url: str = "http://localhost:7860") -> Dict[str, Any]:
+    """Run LLM-based agent on a task."""
+    agent = SQLQualityAgent(base_url)
+
+    if not agent.health_check():
+        print("Error: Environment not available")
+        sys.exit(1)
+
+    # Initialize episode
+    print(f"Initializing task: {task_id}")
+    obs = agent.reset(task_id)
+    print(f"Episode ID: {agent.episode_id}")
+    
+    # Get schema and task info
+    schema = obs.get("observation", {}).get("schema_info", [])
+    task_desc = obs.get("observation", {}).get("task_description", "")
+    
+    print(f"Task: {task_desc}")
+    print(f"Schema: {json.dumps(schema, indent=2)}")
+    
+    # Build context for LLM
+    schema_text = json.dumps(schema, indent=2)
+    
+    try:
+        if task_id == "null_audit":
+            # Use LLM to find null percentages
+            prompt = f"""Given this schema: {schema_text}
+            
+Write a SQL query to find which column has the highest percentage of NULL values.
+Return only the SQL query."""
+            
+            sql = agent.llm_generate_sql(prompt, "You are a data quality expert.")
+            if sql:
+                result = agent.query(sql)
+                print(f"Query result: {result}")
+            
+            # Ask LLM to identify the answer
+            answer_prompt = f"""Based on the query results above, which column name has null rate > 50%?
+Return ONLY the column name (e.g., 'city'), nothing else."""
+            
+            answer = agent.llm_generate_sql(answer_prompt, "You identify the problematic column.")
+            if not answer:
+                answer = "city"  # fallback
+            
+            submit_result = agent.submit(answer=answer.strip())
+            return submit_result
+
+        elif task_id == "type_repair":
+            # Check for type issues
+            prompt = f"""Given schema: {schema_text}
+            
+Task: {task_desc}
+Write a SQL query to inspect the transactions table and identify type mismatches in the amount column.
+Return only the SQL."""
+            
+            sql = agent.llm_generate_sql(prompt)
+            if sql:
+                result = agent.query(sql)
+                print(f"Inspection result: {result}")
+            
+            # Generate fix
+            fix_prompt = f"""Given the type mismatch in the amount column (stored as TEXT with '$' prefix),
+write a SQLite UPDATE statement to convert it to REAL numeric values.
+Return only the SQL UPDATE statement."""
+            
+            fix_sql = agent.llm_generate_sql(fix_prompt, "You write data repair SQL.")
+            if fix_sql:
+                fix_result = agent.fix(fix_sql)
+                print(f"Fix result: {fix_result}")
+                submit_result = agent.submit(fix_query=fix_sql)
+            else:
+                submit_result = agent.submit()
+            return submit_result
+
+        elif task_id == "dedup_bug":
+            # Find duplicates
+            prompt = f"""Given schema: {schema_text}
+            
+Write a SQL query to find duplicate event_id values in the events table.
+Return only the SQL."""
+            
+            sql = agent.llm_generate_sql(prompt)
+            if sql:
+                result = agent.query(sql)
+                print(f"Duplicates found: {result}")
+            
+            # Generate dedup fix
+            fix_prompt = f"""Write a SQLite DELETE statement to remove duplicate rows from the events table,
+keeping only the first occurrence (minimum rowid) for each event_id.
+Return only the SQL DELETE statement."""
+            
+            fix_sql = agent.llm_generate_sql(fix_prompt)
+            if fix_sql:
+                fix_result = agent.fix(fix_sql)
+                print(f"Fix result: {fix_result}")
+                submit_result = agent.submit(
+                    fix_query=fix_sql,
+                    root_cause_table="events",
+                    offending_key="event_id"
+                )
+            else:
+                submit_result = agent.submit()
+            return submit_result
+
+        elif task_id == "fk_integrity":
+            # Find orphans
+            prompt = f"""Given schema with tables: invoices and invoice_items
+Write a SQL query to count orphan invoice_items (items with invoice_id not in invoices table).
+Return only the SQL."""
+            
+            sql = agent.llm_generate_sql(prompt)
+            if sql:
+                result = agent.query(sql)
+                print(f"Orphan count: {result}")
+            
+            # Generate orphan removal
+            fix_prompt = f"""Write a SQLite DELETE statement to remove orphan records from invoice_items table
+where invoice_id does not exist in the invoices table.
+Return only the SQL DELETE statement."""
+            
+            fix_sql = agent.llm_generate_sql(fix_prompt)
+            if fix_sql:
+                fix_result = agent.fix(fix_sql)
+                print(f"Fix result: {fix_result}")
+                submit_result = agent.submit(fix_query=fix_sql)
+            else:
+                submit_result = agent.submit()
+            return submit_result
+
+        elif task_id == "schema_drift":
+            # Check schema
+            prompt = f"""Write a SQL query: PRAGMA table_info(raw_events)
+Return only the SQL."""
+            
+            sql = agent.llm_generate_sql(prompt)
+            if sql:
+                result = agent.query(sql)
+                print(f"Schema: {result}")
+            
+            # Generate normalization
+            fix_prompt = f"""Given the raw_events table with columns: event_id, timestamp, user_id, value, event_type
+Write a SQLite CREATE TABLE AS statement to create analytics_events with:
+- event_id (preserved)
+- timestamp cast to INTEGER
+- user_id (preserved)
+- value cast to REAL
+- event_type (preserved)
+Return only the SQL CREATE TABLE AS statement."""
+            
+            fix_sql = agent.llm_generate_sql(fix_prompt)
+            if fix_sql:
+                fix_result = agent.fix(fix_sql)
+                print(f"Fix result: {fix_result}")
+                submit_result = agent.submit(fix_query=fix_sql)
+            else:
+                submit_result = agent.submit()
+            return submit_result
+
+        else:
+            raise ValueError(f"Unknown task: {task_id}")
+            
+    except RuntimeError as e:
+        print(f"Agent failed: {e}")
+        # Submit empty answer to end episode
+        return agent.submit()
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Run SQL Quality Agent")
+    parser.add_argument("--task", required=True, help="Task ID to run")
+    parser.add_argument("--url", default="http://localhost:7860", help="Base URL")
+    parser.add_argument("--seed", type=int, help="Random seed")
+
+    args = parser.parse_args()
+
+    result = run_agent(args.task, args.url)
+    print(f"\nFinal result: {json.dumps(result, indent=2)}")