Delete inference.py

inference.py

@@ -1,356 +0,0 @@
-"""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)}")