Hariprita/nl2sql-openenv

1import os2import re3import json4import textwrap5from typing import List, Optional6from openai import OpenAI7from client import SQLAgentEnv, SQLAction8 9API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")10API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY")11MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-Coder-7B-Instruct")12ENV_URL = os.getenv("ENV_URL", "http://localhost:7860")13BENCHMARK = "nl2sql"14MAX_STEPS = 2015TEMPERATURE = 0.116MAX_TOKENS = 30017FALLBACK_SQL = "SELECT 1;"18SUCCESS_SCORE_THRESHOLD = 0.519 20SYSTEM_PROMPT = """You are an expert SQL query writer for SQLite databases.21Respond with EXACTLY one valid SQL SELECT query.22No explanation, no markdown, no backticks.23Never write DROP, DELETE, INSERT, UPDATE, CREATE, ALTER, or TRUNCATE."""24 25 26def log_start(task, env, model):27    print(f"[START] task={task} env={env} model={model}", flush=True)28 29 30def log_step(step, action, reward, done, error):31    error_val = error if error else "null"32    done_val = str(done).lower()33    action_clean = action.replace("\n", " ")[:80]34    print(f"[STEP] step={step} action={action_clean} reward={reward:.2f} done={done_val} error={error_val}", flush=True)35 36 37def log_end(success, steps, score, rewards):38    rewards_str = ",".join(f"{r:.2f}" for r in rewards)39    print(f"[END] success={str(success).lower()} steps={steps} score={score:.2f} rewards={rewards_str}", flush=True)40 41 42def extract_sql(text):43    if not text:44        return FALLBACK_SQL45    text = re.sub(r"```sql\s*", "", text, flags=re.IGNORECASE)46    text = re.sub(r"```\s*", "", text)47    text = text.strip()48    if re.match(r"^\s*(SELECT|WITH)\b", text, re.IGNORECASE):49        return text50    match = re.search(r"(SELECT|WITH)\b.*", text, re.IGNORECASE | re.DOTALL)51    if match:52        return match.group(0).strip()53    return FALLBACK_SQL54 55 56def build_prompt(observation, history):57    hint = getattr(observation, "hint", "")58    hint_line = f"Hint: {hint}" if hint else ""59    prev = "\n".join(history[-3:]) if history else "None"60    return f"""Schema:61{observation.schema}62 63Question: {observation.question}64{hint_line}65Task: {observation.task_id} ({observation.task_difficulty}) | Attempt {observation.attempt}/{observation.max_attempts}66Previous queries:67{prev}68 69Write a SQL SELECT query."""70 71 72def ask_llm(client, observation, history):73    try:74        completion = client.chat.completions.create(75            model=MODEL_NAME,76            messages=[77                {"role": "system", "content": SYSTEM_PROMPT},78                {"role": "user", "content": build_prompt(observation, history)},79            ],80            temperature=TEMPERATURE,81            max_tokens=MAX_TOKENS,82        )83        return extract_sql(completion.choices[0].message.content or "")84    except Exception as e:85        print(f"[DEBUG] LLM error: {e}", flush=True)86        return FALLBACK_SQL87 88 89def main():90    llm = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)91    env = SQLAgentEnv(base_url=ENV_URL)92    history = []93    rewards = []94    steps_taken = 095    score = 0.096    success = False97    episode_results = []98 99    log_start(task="nl2sql-3task", env=BENCHMARK, model=MODEL_NAME)100 101    try:102        result = env.reset()103        observation = result.observation104 105        for step in range(1, MAX_STEPS + 1):106            if result.done:107                break108 109            sql = ask_llm(llm, observation, history)110            result = env.step(SQLAction(sql_query=sql))111            observation = result.observation112 113            reward = result.reward114            done = result.done115            error = observation.error_output if hasattr(observation, "error_output") and observation.error_output else None116 117            rewards.append(reward)118            steps_taken = step119 120            log_step(step=step, action=sql, reward=reward, done=done, error=error)121 122            history.append(f"Step {step}: reward={reward:.2f}")123            episode_results.append({124                "step": step,125                "task_id": observation.task_id,126                "reward": reward,127                "sql": sql,128            })129 130            if done:131                break132 133        score = sum(rewards) / 3.0 if rewards else 0.0134        score = min(max(score, 0.0), 1.0)135        success = score >= SUCCESS_SCORE_THRESHOLD136 137    except Exception as e:138        print(f"[DEBUG] Episode error: {e}", flush=True)139    finally:140        try:141            env.close()142        except Exception as e:143            print(f"[DEBUG] env.close() error: {e}", flush=True)144        log_end(success=success, steps=steps_taken, score=score, rewards=rewards)145 146    print(json.dumps(episode_results, indent=2), flush=True)147 148 149if __name__ == "__main__":150    main()151