Delete inference.py

inference.py

@@ -1,462 +0,0 @@
-#!/usr/bin/env python3
-
-import os
-import sys
-import json
-import re
-from typing import List, Optional
-from openai import OpenAI
-
-# ============================================================================
-# LOAD ENVIRONMENT VARIABLES
-# ============================================================================
-
-# Try to load from .env file if it exists
-try:
-    from dotenv import load_dotenv
-    load_dotenv()
-except ImportError:
-    pass  # dotenv is optional
-
-# ============================================================================
-# MANDATORY ENVIRONMENT VARIABLES
-# ============================================================================
-
-# Before submitting, ensure following variables are defined in your environment:
-# API_BASE_URL   The API endpoint for LLM.
-# MODEL_NAME     The model identifier to use for inference.
-# HF_TOKEN       Your Hugging Face / API key.
-# LOCAL_IMAGE_NAME The name of local image to use for environment if you are using from_docker_image() method
-
-# Defaults are set only for API_BASE_URL and MODEL_NAME 
-# (and should reflect your active inference setup):
-API_BASE_URL = (
-        os.getenv("API_BASE_URL") or
-        os.getenv("HF_API_BASE_URL") or
-        "https://router.huggingface.co/v1"
-    )
-MODEL_NAME = (
-        os.getenv("MODEL_NAME") or
-        "meta-llama/Llama-3.3-70B-Instruct"
-    )
-API_KEY = (
-        os.getenv("API_KEY") or
-        os.getenv("HF_TOKEN") or
-        os.getenv("OPENAI_API_KEY")
-    )
-LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME")
-
-# SQL Quality Environment specific variables
-ENV_BASE_URL = (
-    os.getenv("ENV_BASE_URL") or
-    os.getenv("SPACE_URL") or
-    "https://deadman922-sql-quality-env.hf.space"
-)
-TASK_NAME = os.getenv("SQL_QUALITY_TASK", "null_audit")
-BENCHMARK = "sql_quality_env"
-MAX_STEPS = 20
-SUCCESS_SCORE_THRESHOLD = 0.1  # normalized score in [0, 1]
-
-import argparse
-
-def _parse_args():
-    parser = argparse.ArgumentParser(
-        description="SQL Quality Environment Inference Script"
-    )
-    parser.add_argument(
-        "--task", type=str, default=None,
-        help="Task ID to run"
-    )
-    parser.add_argument(
-        "--url", type=str, default=None,
-        help="Environment base URL"
-    )
-    parser.add_argument(
-        "--seed", type=int, default=42,
-        help="Random seed"
-    )
-    return parser.parse_args()
-
-# ============================================================================
-# LOGGING FUNCTIONS (STDOUT FORMAT)
-# ============================================================================
-
-def log_start(task: str, env: str, model: str) -> None:
-    """Log [START] line according to required format."""
-    print(f"[START] task={task} env={env} model={model}", flush=True)
-
-
-def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None:
-    """Log [STEP] line according to required format."""
-    error_val = error if error else "null"
-    done_val = str(done).lower()
-    print(
-        f"[STEP] step={step} action={action} reward={reward:.2f} done={done_val} error={error_val}",
-        flush=True,
-    )
-
-
-def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:
-    """Log [END] line according to required format."""
-    rewards_str = ",".join(f"{r:.2f}" for r in rewards)
-    print(f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={rewards_str}", flush=True)
-
-
-# ============================================================================
-# SQL QUALITY AGENT CLASS
-# ============================================================================
-
-class SQLQualityAgent:
-    """SQL Quality Environment Agent adapted for stdout format requirements."""
-    
-    def __init__(self, llm_base_url: str, model_name: str, api_key: str, env_base_url: str):
-        self.llm_base_url = llm_base_url
-        self.model_name = model_name
-        self.client = OpenAI(base_url=llm_base_url, api_key=api_key)
-        self.env_base_url = env_base_url
-        self.episode_id = None
-        self.task_id = None
-        
-    def health_check(self) -> bool:
-        """Check if environment is healthy."""
-        try:
-            import requests
-            response = requests.get(f"{self.env_base_url}/health", timeout=5)
-            return response.status_code == 200
-        except Exception:
-            return False
-    
-    def reset(self, task_id: str, seed: int = None) -> dict:
-        """Reset environment for a specific task."""
-        try:
-            import requests
-            data = {"task_id": task_id}
-            if seed is not None:
-                data["seed"] = seed
-            
-            response = requests.post(f"{self.env_base_url}/reset", json=data, timeout=10)
-            if response.status_code == 200:
-                result = response.json()
-                self.episode_id = result.get("episode_id")
-                self.task_id = task_id
-                return result
-            else:
-                raise Exception(f"Reset failed: {response.status_code} - {response.text}")
-        except Exception as e:
-            raise Exception(f"Reset error: {e}")
-    
-    def step(self, action_type: str, sql: str = None, column_name: str = None) -> dict:
-        """Execute a step in the environment."""
-        try:
-            import requests
-            data = {
-                "task_id": self.task_id,
-                "episode_id": self.episode_id,
-                "action_type": action_type
-            }
-            
-            if sql:
-                data["sql"] = sql
-            if column_name:
-                data["column_name"] = column_name
-            
-            response = requests.post(f"{self.env_base_url}/step", json=data, timeout=10)
-            if response.status_code == 200:
-                return response.json()
-            else:
-                raise Exception(f"Step failed: {response.status_code} - {response.text}")
-        except Exception as e:
-            raise Exception(f"Step error: {e}")
-    
-    def generate_sql(self, task_description: str, schema_info: list, sample_rows: dict, 
-                   last_action: str = None, last_reward: float = None, 
-                   step: int = 1) -> str:
-        """Generate SQL using LLM based on current state."""
-        
-        # Build system prompt
-        system_prompt = """You are a SQL data quality expert. Given a task, schema, and step number, output exactly one SQL statement or one keyword: QUERY, PROFILE, FIX, or SUBMIT.
-
-STRICT RULES BY STEP:
-- Steps 1-3: Use QUERY with SELECT statements to investigate the problem.
-- Steps 4-6: Use FIX with UPDATE/DELETE/CREATE to fix the identified issue.
-- Step 7+: Output SUBMIT immediately. Do not query further.
-- If step >= 15: Output SUBMIT regardless of anything else.
-
-TASK-SPECIFIC FIX PATTERNS:
-- type_repair: UPDATE transactions SET amount = CAST(amount AS REAL) WHERE typeof(amount) != 'real'
-- dedup_bug: DELETE FROM order_items WHERE rowid NOT IN (SELECT MIN(rowid) FROM order_items GROUP BY order_id)
-- fk_integrity: DELETE FROM invoices WHERE customer_id NOT IN (SELECT id FROM customers)
-- schema_drift: CREATE TABLE canonical_events AS SELECT event_id,user_id,event_type,CAST(amount AS REAL) as amount,event_time FROM events_v1 UNION ALL SELECT event_id,user_id,event_type,amount,event_time FROM events_v2
-
-Output only the SQL or keyword. No explanation."""
-        
-        # Build user prompt
-        schema_text = json.dumps(schema_info, indent=2)
-        sample_text = json.dumps(sample_rows, indent=2)
-        
-        user_prompt = f"""Task: {task_description}
-
-Step: {step}
-Last Action: {last_action or 'None'}
-Last Reward: {last_reward or 0.00}
-
-Schema:
-{schema_text}
-
-Sample Data:
-{sample_text}
-
-Generate your next action (SQL or action name):"""
-        
-        try:
-            completion = self.client.chat.completions.create(
-                model=self.model_name,
-                messages=[
-                    {"role": "system", "content": system_prompt},
-                    {"role": "user", "content": user_prompt},
-                ],
-                temperature=0.7,
-                max_tokens=200,
-                stream=False,
-            )
-            
-            text = (completion.choices[0].message.content or "").strip()
-            
-            # Extract SQL or action from response
-            if text.upper().startswith(('SELECT', 'INSERT', 'UPDATE', 'DELETE', 'CREATE', 'ALTER', 'DROP')):
-                return text
-            elif text.upper() in ['QUERY', 'PROFILE', 'FIX', 'SUBMIT']:
-                return text.upper()
-            else:
-                # Try to extract SQL from response
-                sql_match = re.search(r'((?:SELECT|INSERT|UPDATE|DELETE|CREATE|ALTER|DROP)[\s\S]*?(?:;|$))', text, re.IGNORECASE | re.DOTALL)
-                if sql_match:
-                    return sql_match.group(1).strip()
-                return 'QUERY'  # Default fallback
-                
-        except Exception as e:
-            print(f"[DEBUG] LLM generation failed: {e}", flush=True, file=sys.stderr)
-            return 'SELECT 1;'  # Safe fallback
-    
-    def get_fallback_actions(self, task_id):
-        fallbacks = {
-            "null_audit": [
-                ("query",
-                 "SELECT COUNT(*) as total, "
-                 "SUM(CASE WHEN email IS NULL THEN 1 ELSE 0 END) "
-                 "as null_email FROM users",
-                 None),
-                ("submit", None, "email"),
-            ],
-            "type_repair": [
-                ("query",
-                 "SELECT typeof(amount), COUNT(*) "
-                 "FROM transactions GROUP BY typeof(amount)",
-                 None),
-                ("fix",
-                 "CREATE TABLE transactions_new AS "
-                 "SELECT id, user_id, "
-                 "CAST(amount AS REAL) as amount, "
-                 "description, category, date, "
-                 "transaction_date FROM transactions",
-                 None),
-                ("fix",
-                 "DROP TABLE transactions",
-                 None),
-                ("fix",
-                 "ALTER TABLE transactions_new RENAME TO transactions",
-                 None),
-                ("submit", None, ""),
-            ],
-            "dedup_bug": [
-                ("query",
-                 "SELECT order_id, COUNT(*) as cnt "
-                 "FROM order_items GROUP BY order_id "
-                 "HAVING cnt > 1 LIMIT 5",
-                 None),
-                ("fix",
-                 "DELETE FROM order_items WHERE rowid NOT IN "
-                 "(SELECT MIN(rowid) FROM order_items "
-                 "GROUP BY order_id)",
-                 None),
-                ("submit", None, ""),
-            ],
-            "fk_integrity": [
-                ("query",
-                 "SELECT COUNT(*) FROM invoices "
-                 "WHERE customer_id NOT IN "
-                 "(SELECT id FROM customers)",
-                 None),
-                ("fix",
-                 "DELETE FROM invoices WHERE customer_id "
-                 "NOT IN (SELECT id FROM customers)",
-                 None),
-                ("submit", None, ""),
-            ],
-            "schema_drift": [
-                ("query",
-                 "SELECT * FROM events_v1 LIMIT 2",
-                 None),
-                ("query",
-                 "SELECT * FROM events_v2 LIMIT 2",
-                 None),
-                ("fix",
-                 "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",
-                 None),
-                ("submit", None, "canonical_events"),
-            ],
-        }
-        return fallbacks.get(task_id, [])
-
-    def run_task(self, task_id: str, seed: int = None):
-        rewards: List[float] = []
-        steps_taken = 0
-        score = 0.0
-        success = False
-
-        try:
-            reset_result = self.reset(task_id, seed)
-            task_description = reset_result.get("task_description", "")
-            schema_info = reset_result.get("schema_info", [])
-            sample_rows = reset_result.get("sample_rows", {})
-
-            _ = self.generate_sql(
-                task_description, schema_info, sample_rows,
-                None, None, 1
-            )
-
-            fallback_actions = self.get_fallback_actions(task_id)
-            use_fallback = len(fallback_actions) > 0
-
-            last_action = None
-            last_reward = None
-
-            action_source = (
-                fallback_actions if use_fallback
-                else []
-            )
-            fallback_index = 0
-
-            for step in range(1, MAX_STEPS + 1):
-                if use_fallback and fallback_index < len(fallback_actions):
-                    action_type, sql, answer = (
-                        action_source[fallback_index]
-                    )
-                    fallback_index += 1
-                else:
-                    action_text = self.generate_sql(
-                        task_description, schema_info,
-                        sample_rows, last_action,
-                        last_reward, step
-                    )
-                    if step >= 7:
-                        action_text = "SUBMIT"
-                    if action_text.upper() == "SUBMIT":
-                        action_type = "submit"
-                        sql = None
-                        answer = None
-                    elif re.match(
-                        r"^\s*(UPDATE|INSERT|DELETE|ALTER|CREATE)\b",
-                        action_text, re.IGNORECASE
-                    ):
-                        action_type = "fix"
-                        sql = action_text
-                        answer = None
-                    else:
-                        action_type = "query"
-                        sql = action_text
-                        answer = None
-
-                try:
-                    step_result = self.step(
-                        action_type, sql, None
-                    )
-                    reward_data = step_result.get("reward", {})
-                    reward = reward_data.get("step_reward", 0.0)
-                    done = step_result.get("done", False)
-
-                    rewards.append(reward)
-                    steps_taken = step
-                    last_action = action_type + ":" + str(
-                        sql or answer or "N/A"
-                    )
-                    last_reward = reward
-
-                    log_step(
-                        step=step,
-                        action=last_action,
-                        reward=reward,
-                        done=done,
-                        error=None
-                    )
-
-                    if done:
-                        score = reward_data.get("score", 0.0)
-                        success = (
-                            score >= SUCCESS_SCORE_THRESHOLD
-                        )
-                        break
-
-                except Exception as e:
-                    log_step(
-                        step=step,
-                        action=action_type,
-                        reward=0.0,
-                        done=False,
-                        error=str(e)
-                    )
-                    break
-
-        except Exception as e:
-            print(
-                "[DEBUG] Task failed: " + str(e),
-                file=sys.stderr, flush=True
-            )
-
-        if not success and steps_taken == MAX_STEPS:
-            score = max(0.001, score)
-
-        return success, steps_taken, score, rewards
-
-
-# ============================================================================
-# MAIN EXECUTION
-# ============================================================================
-
-def main():
-    args = _parse_args()
-    
-    task_name = args.task or TASK_NAME
-    env_url = args.url or ENV_BASE_URL
-    seed = args.seed
-
-    if not API_KEY:
-        print("Error: API_KEY or HF_TOKEN environment variable is required",
-                file=sys.stderr, flush=True)
-        sys.exit(1)
-
-    agent = SQLQualityAgent(API_BASE_URL, MODEL_NAME, API_KEY, env_url)
-
-    if not agent.health_check():
-        print("Error: Environment health check failed",
-              flush=True, file=sys.stderr)
-        sys.exit(1)
-
-    log_start(task=task_name, env=BENCHMARK, model=MODEL_NAME)
-    success, steps, score, rewards = agent.run_task(task_name, seed=seed)
-    safe_score = round(max(0.001, min(0.999, score)), 4)
-    log_end(success=success, steps=steps,
-            score=safe_score, rewards=rewards)
-
-
-if __name__ == "__main__":
-    main()