SyncShift/sql-correction-env

1"""2inference.py - SQL Correction Environment Baseline Script3=========================================================4MANDATORY - Place this file in the ROOT of the project.5 6Required environment variables:7  API_BASE_URL   The API endpoint for the LLM8  MODEL_NAME     The model identifier to use for inference9  HF_TOKEN       Your Hugging Face / API key10  ENV_URL        URL of the running environment (default: http://localhost:7860)11  SQL_ENV_TASK   Task difficulty: easy | medium | hard (default: easy)12"""13 14import asyncio15import os16import sys17import textwrap18from typing import List, Optional19import re20 21import httpx22 23try:24    from openai import OpenAI25except Exception:26    OpenAI = None27 28API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")29MODEL_NAME   = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")30API_KEY      = os.getenv("HF_TOKEN") or os.getenv("API_KEY", "dummy")31TASK_NAME    = os.getenv("SQL_ENV_TASK", "easy")32BENCHMARK    = "sql-correction-env"33ENV_URL      = os.getenv("ENV_URL", "http://localhost:7860")34MAX_STEPS    = 835SUCCESS_SCORE_THRESHOLD = 0.536 37 38# ---------------------------------------------------------------------------39# Logging helpers — must match the spec format exactly40# ---------------------------------------------------------------------------41 42def log_start(task: str, env: str, model: str) -> None:43    print(f"[START] task={task} env={env} model={model}", flush=True)44 45 46def log_step(47    step: int,48    action: str,49    reward: float,50    done: bool,51    error: Optional[str],52) -> None:53    err = error if error else "null"54    done_val = str(done).lower()55    action_clean = action.replace("\n", " ").replace("\r", "").strip()56    print(57        f"[STEP] step={step} action={action_clean} "58        f"reward={reward:.2f} done={done_val} error={err}",59        flush=True,60    )61 62 63def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:64    rewards_str = ",".join(f"{r:.2f}" for r in rewards) if rewards else ""65    print(66        f"[END] success={str(success).lower()} steps={steps} "67        f"score={score:.3f} rewards={rewards_str}",68        flush=True,69    )70 71 72# ---------------------------------------------------------------------------73# LLM / heuristic helpers74# ---------------------------------------------------------------------------75 76SYSTEM_PROMPT = textwrap.dedent(77    """78    You are an expert SQL debugger.79    You will be shown a broken SQL query that contains typos or keyword errors.80    Fix ALL errors and return ONLY the corrected SQL query.81    No explanation, no markdown, no code blocks, no backticks.82    Common keyword typos to watch for:83      FORM->FROM, WEHRE->WHERE, WHER->WHERE,84      GRUP->GROUP, HAVNG->HAVING, ORDR->ORDER,85      INNE->INNER, LFT->LEFT, BETWEN->BETWEEN,86      DSC->DESC, SELCT->SELECT, LIMT->LIMIT.87    Also watch for column name errors described in the schema context.88    """89).strip()90 91SQL_REPLACEMENTS = {92    "FORM":   "FROM",93    "WEHRE":  "WHERE",94    "WHER":   "WHERE",95    "GRUP":   "GROUP",96    "HAVNG":  "HAVING",97    "ORDR":   "ORDER",98    "INNE":   "INNER",99    "LFT":    "LEFT",100    "BETWEN": "BETWEEN",101    "DSC":    "DESC",102    "SELCT":  "SELECT",103    "LIMT":   "LIMIT",104}105 106 107def heuristic_correct_sql(query: str) -> str:108    """Deterministic fallback when the LLM is unavailable."""109    corrected = query110    for broken, fixed in SQL_REPLACEMENTS.items():111        corrected = re.sub(112            rf"\b{re.escape(broken)}\b", fixed, corrected, flags=re.IGNORECASE113        )114    return corrected.strip()115 116 117def get_model_action(118    client: Optional["OpenAI"],119    obs: dict,120    history: List[str],121) -> str:122    """Return a corrected SQL string. Falls back to heuristic on any failure."""123    heuristic = heuristic_correct_sql(obs.get("broken_query", ""))124 125    if client is None:126        return heuristic127 128    history_block = "\n".join(history[-4:]) if history else "None"129    user_prompt = textwrap.dedent(130        f"""131        Broken SQL query:132        {obs.get("broken_query", "")}133 134        Schema context:    {obs.get("schema_context") or "Not provided"}135        Error hint:        {obs.get("error_hint") or "None"}136        Steps remaining:   {obs.get("steps_remaining", "?")}137        Previous attempt:  {obs.get("previous_attempt") or "None"}138        Feedback:          {obs.get("feedback") or "None"}139 140        Recent history:141        {history_block}142 143        Return ONLY the corrected SQL query.144        """145    ).strip()146 147    try:148        completion = client.chat.completions.create(149            model=MODEL_NAME,150            messages=[151                {"role": "system", "content": SYSTEM_PROMPT},152                {"role": "user",   "content": user_prompt},153            ],154            temperature=0.2,155            max_tokens=300,156            stream=False,157        )158        text = (completion.choices[0].message.content or "").strip()159        return text if text else heuristic160    except Exception as exc:161        print(f"[DEBUG] LLM call failed: {exc}", flush=True)162        return heuristic163 164 165# ---------------------------------------------------------------------------166# Episode runner167# ---------------------------------------------------------------------------168 169async def run_task(task_name: str) -> None:170    """171    Run one full episode for `task_name`.172 173    The [END] log line is ALWAYS emitted via the finally block, even if an174    exception occurs mid-episode or the reset call fails.175    """176    rewards: List[float] = []177    history: List[str]   = []178    steps_taken = 0179    score       = 0.0180    success     = False181 182    client = None183    if OpenAI is not None and API_KEY not in {"", "dummy"}:184        try:185            client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)186        except Exception as exc:187            print(f"[DEBUG] OpenAI client init failed: {exc}", flush=True)188 189    log_start(task_name, BENCHMARK, MODEL_NAME)190 191    http: Optional[httpx.AsyncClient] = None192    try:193        http = httpx.AsyncClient(base_url=ENV_URL, timeout=60.0)194 195        # ── reset ──────────────────────────────────────────────────────────196        reset_failed = False197        obs: dict = {}198        try:199            reset_resp = await http.post(200                "/reset", json={"difficulty": task_name}201            )202            reset_resp.raise_for_status()203            reset_data = reset_resp.json()204            # The openenv wrapper may nest the observation under "observation"205            obs = reset_data.get("observation", reset_data)206        except Exception as exc:207            print(f"[DEBUG] Reset failed: {exc}", flush=True)208            reset_failed = True209 210        if not reset_failed:211            # ── step loop ──────────────────────────────────────────────────212            for step in range(1, MAX_STEPS + 1):213                try:214                    action_str = get_model_action(client, obs, history)215                except Exception as exc:216                    print(f"[DEBUG] Model action failed: {exc}", flush=True)217                    action_str = heuristic_correct_sql(218                        obs.get("broken_query", "")219                    )220 221                try:222                    # Action must be wrapped under {"action": {...}}223                    step_resp = await http.post(224                        "/step",225                        json={"action": {"corrected_query": action_str}},226                    )227                    step_resp.raise_for_status()228                    result = step_resp.json()229                except Exception as exc:230                    print(f"[DEBUG] Step {step} request failed: {exc}", flush=True)231                    rewards.append(0.0)232                    steps_taken = step233                    log_step(step, action_str, 0.0, True, str(exc))234                    break235 236                obs    = result.get("observation", obs)237                reward = float(result.get("reward", 0.0))238                done   = bool(result.get("done", False))239                info   = result.get("info")240                error  = info.get("error") if isinstance(info, dict) else None241 242                rewards.append(reward)243                steps_taken = step244                history.append(245                    f"Step {step}: attempt={action_str!r} reward={reward:+.2f}"246                )247 248                log_step(step, action_str, reward, done, error)249 250                if done:251                    break252 253            if rewards:254                score   = min(max(sum(rewards) / len(rewards), 0.01), 0.99)255            success = score >= SUCCESS_SCORE_THRESHOLD256 257    except Exception as exc:258        print(f"[DEBUG] Unhandled episode error: {exc}", flush=True)259 260    finally:261        if http is not None:262            try:263                await http.aclose()264            except Exception as exc:265                print(f"[DEBUG] HTTP close error: {exc}", flush=True)266        log_end(success, steps_taken, score, rewards)267 268 269# ---------------------------------------------------------------------------270# Entry point271# ---------------------------------------------------------------------------272 273async def main() -> None:274    """Run all three difficulties in sequence so validator sees 3 [END] lines."""275    try:276        for difficulty in ("easy", "medium", "hard"):277            await run_task(difficulty)278            print("", flush=True)279    except Exception as exc:280        print(f"[DEBUG] Main loop error: {exc}", flush=True)281 282 283if __name__ == "__main__":284    try:285        asyncio.run(main())286    except KeyboardInterrupt:287        pass288    except Exception as exc:289        print(f"[DEBUG] Fatal error: {exc}", flush=True)290    finally:291        sys.exit(0)292