VarshaChukka/SQL_Debugging_REPL_Environment

1import os2from openai import OpenAI3import requests4 5# REQUIRED ENV VARIABLES6API_BASE_URL = os.getenv("API_BASE_URL", "http://127.0.0.1:8000")7MODEL_NAME = os.getenv("MODEL_NAME", "gpt-4.1-mini")8HF_TOKEN = os.getenv("HF_TOKEN")9 10if HF_TOKEN is None:11    raise ValueError("HF_TOKEN environment variable is required")12 13# Initialize OpenAI client (MANDATORY)14client = OpenAI(15    base_url="https://api.openai.com/v1",16    api_key=HF_TOKEN17)18 19BASE_URL = API_BASE_URL20 21 22def generate_sql(problem):23    problem = problem.lower()24 25    if "all users" in problem:26        return "SELECT * FROM users"27 28    if "id > 1" in problem:29        return "SELECT * FROM users WHERE id > 1"30 31    if "join" in problem:32        return "SELECT users.id, users.name, orders.amount FROM users JOIN orders ON users.id = orders.user_id"33 34    return "SELECT * FROM users"35 36 37def run_task(task_id):38    step_count = 039    rewards = []40    success = False41 42    print(f"[START] task={task_id} env=sql-debug model={MODEL_NAME}")43 44    try:45        #FIX: POST instead of GET46        res = requests.post(f"{BASE_URL}/reset?task_id={task_id}")47        res_json = res.json()48        problem = res_json["problem"]49 50        # GENERATE ACTION51        action = generate_sql(problem)52        step_count += 153 54        # STEP55        step_res = requests.post(56            f"{BASE_URL}/step",57            json={"sql_query": action}58        )59        step_json = step_res.json()60 61        reward = float(step_json[1]["score"])62        done = step_json[2]63 64        rewards.append(f"{reward:.2f}")65 66        #FIX: error=null (no quotes)67        print(68            f"[STEP] step={step_count} action={action} reward={reward:.2f} "69            f"done={'true' if done else 'false'} error=null"70        )71 72        success = done73 74    except Exception as e:75        error_msg = str(e).replace("\n", " ")76 77        print(78            f"[STEP] step={step_count} action=null reward=0.00 done=false error=\"{error_msg}\""79        )80 81    print(82        f"[END] success={'true' if success else 'false'} "83        f"steps={step_count} rewards={','.join(rewards) if rewards else '0.00'}"84    )85 86 87if __name__ == "__main__":88    for task_id in range(3):89        run_task(task_id)