Delete inference.py

inference.py

@@ -1,501 +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-compatible API (Groq/HuggingFace) to generate SQL queries and fixes dynamically.
-
-MANDATORY:
-- Before submitting, ensure the following variables are defined in your environment configuration:
-  API_BASE_URL The API endpoint for the LLM.
-  MODEL_NAME  The model identifier to use for inference.
-  HF_TOKEN Your Hugging Face / API key.
-- The inference script must be named `inference.py` and placed in the root directory of the project
-- Participants must use OpenAI Client for all LLM calls using above variables
-"""
-
-import json
-import os
-import sys
-from typing import Any, Dict, List, Optional
-
-import requests
-from openai import OpenAI
-try:
-    from dotenv import load_dotenv
-except Exception:  # pragma: no cover - optional dependency guard
-    load_dotenv = None
-
-if load_dotenv:
-    load_dotenv()
-
-# Configuration from environment
-MAX_STEPS = int(os.getenv("MAX_STEPS", "10"))
-API_BASE_URL = os.getenv("API_BASE_URL", "<your-active-endpoint>")
-MODEL_NAME = os.getenv("MODEL_NAME", "<your-active-model>")
-HF_TOKEN = os.getenv("HF_TOKEN")
-API_KEY = os.getenv("GROQ_API_KEY", "") or HF_TOKEN or os.getenv("OPENAI_API_KEY", "")
-
-
-# Structured logging utilities for validator parsing
-def log_start(episode_id: str, task_id: str, timestamp: float = None):
-    ts = timestamp or __import__('time').time()
-    print(f"START episode_id={episode_id} task_id={task_id} timestamp={ts}")
-
-
-def log_step(step: int, action: str, reward: float, score: float = None, sql: str = None, column: str = None):
-    extras = []
-    if sql:
-        extras.append(f'sql="{sql[:50]}..."' if len(sql) > 50 else f'sql="{sql}"')
-    if column:
-        extras.append(f'column="{column}"')
-    if score is not None:
-        extras.append(f'score={score:.4f}')
-    extra_str = " " + " ".join(extras) if extras else ""
-    print(f"STEP step={step} action={action} reward={reward:.4f}{extra_str}")
-
-
-def log_end(episode_id: str, final_score: float, total_steps: int):
-    print(f"END episode_id={episode_id} final_score={final_score:.4f} total_steps={total_steps}")
-
-
-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-compatible client (works with Groq, HF, OpenAI)
-        self.client = OpenAI(
-            api_key=API_KEY,
-            base_url=API_BASE_URL
-        )
-
-    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
-        self.step_count = 0
-        self.trajectory = []
-        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()
-        result = r.json()
-        self.trajectory.append({"type": "query", "sql": sql, "result": result})
-        return result
-
-    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()
-        result = r.json()
-        self.trajectory.append({"type": "profile", "column": column, "result": result})
-        return result
-
-    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()
-        result = r.json()
-        self.trajectory.append({"type": "fix", "sql": fix_sql, "result": result})
-        return result
-
-    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(self, prompt: str, system_context: str = "") -> str:
-        """Generate text using LLM via OpenAI-compatible API."""
-        messages = [
-            {"role": "system", "content": system_context or "You are a SQL data quality expert. Generate valid SQLite SQL queries. Reply with ONLY the SQL, no explanations."},
-            {"role": "user", "content": prompt}
-        ]
-        
-        try:
-            response = self.client.chat.completions.create(
-                model=MODEL_NAME,
-                messages=messages,
-                temperature=0.1,
-                max_tokens=500
-            )
-            text = response.choices[0].message.content.strip()
-            # Extract SQL from markdown if present
-            if "```sql" in text:
-                text = text.split("```sql")[1].split("```")[0].strip()
-            elif "```" in text:
-                text = text.split("```")[1].split("```")[0].strip()
-            return text
-        except Exception as e:
-            print(f"LLM generation failed: {e}")
-            return ""
-
-    def llm_answer(self, prompt: str, system_context: str = "") -> str:
-        """Generate a short answer using LLM."""
-        messages = [
-            {"role": "system", "content": system_context or "You are a data quality expert. Reply with ONLY the answer, no explanations."},
-            {"role": "user", "content": prompt}
-        ]
-        
-        try:
-            response = self.client.chat.completions.create(
-                model=MODEL_NAME,
-                messages=messages,
-                temperature=0.0,
-                max_tokens=50
-            )
-            return response.choices[0].message.content.strip().strip('"').strip("'")
-        except Exception as e:
-            print(f"LLM answer 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 with multi-step investigation."""
-    import time
-    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}")
-    log_start(agent.episode_id, task_id, time.time())
-    
-    # Get schema and task info
-    schema = obs.get("schema_info", [])
-    task_desc = obs.get("task_description", "")
-    
-    print(f"Task: {task_desc}")
-    print(f"Schema: {json.dumps(schema, indent=2)}")
-    
-    schema_text = json.dumps(schema, indent=2)
-    
-    try:
-        if task_id == "null_audit":
-            # Step 1: Query null percentages for each column
-            sql = agent.llm_generate(
-                f"""Given this schema: {schema_text}
-
-Write a SQL query that counts NULL values for each column in the 'users' table 
-and calculates the null percentage. Return column_name, null_count, and pct_null.
-Return only the SQL query."""
-            )
-            if sql:
-                result = agent.query(sql)
-                reward = result.get('reward', {}).get('step_reward', 0.0)
-                log_step(1, "query", reward, sql=sql)
-                print(f"Null analysis: {json.dumps(result.get('reward', {}), indent=2)}")
-            
-            # Step 2: Profile individual columns
-            step_num = 2
-            for col in ["city", "email", "age", "signup_date"]:
-                profile_result = agent.profile(col)
-                reward = profile_result.get('reward', {}).get('step_reward', 0.0)
-                log_step(step_num, "profile", reward, column=col)
-                step_num += 1
-                print(f"Profile {col}: {profile_result.get('reward', {}).get('feedback', '')}")
-            
-            # Step 3: Ask LLM to identify the answer based on profile results
-            trajectory_summary = "\n".join([
-                f"Step {i+1}: {t.get('type', 'unknown')} - {t.get('result', {}).get('reward', {}).get('feedback', '')[:100]}"
-                for i, t in enumerate(agent.trajectory)
-            ])
-            
-            answer = agent.llm_answer(
-                f"""Based on these investigation results:
-{trajectory_summary}
-
-Which column name in the 'users' table has a null rate exceeding 50%?
-Return ONLY the column name (e.g., 'city'), nothing else.""",
-                "You identify the problematic column. Return ONLY the column name."
-            )
-            if not answer:
-                answer = "city"  # fallback
-            
-            print(f"Identified column: {answer}")
-            submit_result = agent.submit(answer=answer.strip())
-            final_score = submit_result.get('reward', {}).get('score', 0.0)
-            log_step(agent.step_count, "submit", final_score, score=final_score)
-            log_end(agent.episode_id, final_score, agent.step_count)
-            return submit_result
-
-        elif task_id == "type_repair":
-            # Step 1: Inspect the table
-            result = agent.query("SELECT typeof(amount), amount FROM transactions LIMIT 10")
-            print(f"Type inspection: {result.get('reward', {}).get('feedback', '')[:100]}")
-            
-            # Step 2: Generate fix using LLM
-            fix_sql = agent.llm_generate(
-                f"""The 'transactions' table has an 'amount' column stored as TEXT with numeric values like '12.50'.
-Write a SQLite script to:
-1. Create a new table with amount as REAL  
-2. Copy all data, casting amount to REAL
-3. Drop the old table  
-4. Rename the new table
-
-Return only the SQL statements."""
-            )
-            
-            if fix_sql:
-                print(f"Generated fix: {fix_sql[:200]}")
-                fix_result = agent.fix(fix_sql)
-                print(f"Fix preview: {fix_result.get('reward', {}).get('feedback', '')}")
-                
-                # Step 3: Verify the fix
-                verify = agent.query("SELECT typeof(amount), amount FROM transactions LIMIT 5")
-                print(f"Verification: {verify.get('reward', {}).get('feedback', '')[:100]}")
-                
-                submit_result = agent.submit(fix_query=fix_sql)
-                final_score = submit_result.get('reward', {}).get('score', 0.0)
-                log_step(agent.step_count, "submit", final_score, score=final_score)
-                log_end(agent.episode_id, final_score, agent.step_count)
-                return submit_result
-            else:
-                # Fallback fix
-                fallback = """
-CREATE TABLE transactions_fixed AS 
-SELECT id, user_id, CAST(amount AS REAL) AS amount, description, category, date, transaction_date 
-FROM transactions;
-DROP TABLE transactions;
-ALTER TABLE transactions_fixed RENAME TO transactions;
-"""
-                agent.fix(fallback)
-                submit_result = agent.submit(fix_query=fallback)
-                final_score = submit_result.get('reward', {}).get('score', 0.0)
-                log_step(agent.step_count, "submit", final_score, score=final_score)
-                log_end(agent.episode_id, final_score, agent.step_count)
-                return submit_result
-
-        elif task_id == "dedup_bug":
-            # Step 1: Count duplicates
-            result = agent.query(
-                "SELECT order_id, COUNT(*) as cnt FROM order_items GROUP BY order_id HAVING cnt > 1 LIMIT 10"
-            )
-            print(f"Duplicates: {result.get('reward', {}).get('feedback', '')[:100]}")
-            
-            # Step 2: Generate dedup fix
-            fix_sql = agent.llm_generate(
-                f"""The 'order_items' table has duplicate rows for the same order_id.
-Write a SQLite DELETE statement to remove duplicate rows, keeping only the row 
-with the smallest rowid for each order_id.
-Return only the SQL DELETE statement."""
-            )
-            
-            if fix_sql:
-                fix_result = agent.fix(fix_sql)
-                print(f"Fix preview: {fix_result.get('reward', {}).get('feedback', '')}")
-                
-                # Verify
-                agent.query("SELECT order_id, COUNT(*) as cnt FROM order_items GROUP BY order_id HAVING cnt > 1")
-                
-                submit_result = agent.submit(
-                    fix_query=fix_sql,
-                    root_cause_table="order_items",
-                    offending_key="order_id"
-                )
-                final_score = submit_result.get('reward', {}).get('score', 0.0)
-                log_step(agent.step_count, "submit", final_score, score=final_score)
-                log_end(agent.episode_id, final_score, agent.step_count)
-                return submit_result
-            else:
-                fallback = "DELETE FROM order_items WHERE rowid NOT IN (SELECT MIN(rowid) FROM order_items GROUP BY order_id)"
-                agent.fix(fallback)
-                submit_result = agent.submit(
-                    fix_query=fallback,
-                    root_cause_table="order_items",
-                    offending_key="order_id"
-                )
-                final_score = submit_result.get('reward', {}).get('score', 0.0)
-                log_step(agent.step_count, "submit", final_score, score=final_score)
-                log_end(agent.episode_id, final_score, agent.step_count)
-                return submit_result
-
-        elif task_id == "fk_integrity":
-            # Step 1: Find orphans
-            result = agent.query(
-                "SELECT COUNT(*) FROM invoices i LEFT JOIN customers c ON i.customer_id = c.id WHERE c.id IS NULL"
-            )
-            print(f"Orphans: {result.get('reward', {}).get('feedback', '')[:100]}")
-            
-            # Step 2: Generate fix
-            fix_sql = agent.llm_generate(
-                f"""The 'invoices' table has orphan records where customer_id doesn't match any id in the 'customers' table.
-Write a SQLite DELETE statement to remove these orphan invoices.
-Return only the SQL DELETE statement."""
-            )
-            
-            if fix_sql:
-                agent.fix(fix_sql)
-                agent.query("SELECT COUNT(*) FROM invoices i LEFT JOIN customers c ON i.customer_id = c.id WHERE c.id IS NULL")
-                submit_result = agent.submit(fix_query=fix_sql)
-                final_score = submit_result.get('reward', {}).get('score', 0.0)
-                log_step(agent.step_count, "submit", final_score, score=final_score)
-                log_end(agent.episode_id, final_score, agent.step_count)
-                return submit_result
-            else:
-                fallback = "DELETE FROM invoices WHERE customer_id NOT IN (SELECT id FROM customers)"
-                agent.fix(fallback)
-                submit_result = agent.submit(fix_query=fallback)
-                final_score = submit_result.get('reward', {}).get('score', 0.0)
-                log_step(agent.step_count, "submit", final_score, score=final_score)
-                log_end(agent.episode_id, final_score, agent.step_count)
-                return submit_result
-
-        elif task_id == "schema_drift":
-            # Step 1: Inspect both tables
-            agent.query("PRAGMA table_info(events_v1)")
-            agent.query("PRAGMA table_info(events_v2)")
-            
-            # Step 2: Generate normalization
-            fix_sql = agent.llm_generate(
-                f"""Two tables exist:
-- events_v1: id, user_id, event_name, amount_text (TEXT), event_time
-- events_v2: event_id, account_id, event_type, amount (REAL), event_ts, source
-
-Create a canonical_events table with columns: event_id, user_id, event_type, amount (REAL), event_time
-by merging both tables. Map:
-  events_v1: id→event_id, user_id→user_id, event_name→event_type, CAST(amount_text AS REAL)→amount, event_time→event_time
-  events_v2: event_id→event_id, account_id→user_id, event_type→event_type, amount→amount, event_ts→event_time
-
-Return only the CREATE TABLE AS and INSERT statements."""
-            )
-            
-            if fix_sql:
-                agent.fix(fix_sql)
-                agent.query("SELECT COUNT(*) FROM canonical_events")
-                submit_result = agent.submit(fix_query=fix_sql)
-                final_score = submit_result.get('reward', {}).get('score', 0.0)
-                log_step(agent.step_count, "submit", final_score, score=final_score)
-                log_end(agent.episode_id, final_score, agent.step_count)
-                return submit_result
-            else:
-                fallback = """
-CREATE TABLE canonical_events AS
-SELECT id AS event_id, user_id, event_name AS event_type, CAST(amount_text AS REAL) AS amount, event_time
-FROM events_v1
-UNION ALL
-SELECT event_id, account_id AS user_id, event_type, amount, event_ts AS event_time
-FROM events_v2;
-"""
-                agent.fix(fallback)
-                submit_result = agent.submit(fix_query=fallback)
-                final_score = submit_result.get('reward', {}).get('score', 0.0)
-                log_step(agent.step_count, "submit", final_score, score=final_score)
-                log_end(agent.episode_id, final_score, agent.step_count)
-                return submit_result
-
-        else:
-            raise ValueError(f"Unknown task: {task_id}")
-            
-    except RuntimeError as e:
-        print(f"Agent failed: {e}")
-        result = agent.submit()
-        final_score = result.get('reward', {}).get('score', 0.0)
-        log_end(agent.episode_id, final_score, agent.step_count)
-        return result
-
-
-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)}")