Spaces:
Sleeping
Sleeping
| """ | |
| inference.py - SQL Correction Environment Baseline Script | |
| ========================================================= | |
| MANDATORY - Place this file in the ROOT of the project. | |
| Required environment variables: | |
| 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 | |
| ENV_URL URL of the running environment (default: http://localhost:7860) | |
| SQL_ENV_TASK Task difficulty: easy | medium | hard (default: easy) | |
| """ | |
| import asyncio | |
| import os | |
| import sys | |
| import textwrap | |
| from typing import List, Optional | |
| import re | |
| import httpx | |
| try: | |
| from openai import OpenAI | |
| except Exception: | |
| OpenAI = None | |
| API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1") | |
| MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct") | |
| API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY", "dummy") | |
| TASK_NAME = os.getenv("SQL_ENV_TASK", "easy") | |
| BENCHMARK = "sql-correction-env" | |
| ENV_URL = os.getenv("ENV_URL", "http://localhost:7860") | |
| MAX_STEPS = 8 | |
| SUCCESS_SCORE_THRESHOLD = 0.5 | |
| # --------------------------------------------------------------------------- | |
| # Logging helpers — must match the spec format exactly | |
| # --------------------------------------------------------------------------- | |
| def log_start(task: str, env: str, model: str) -> None: | |
| 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: | |
| err = error if error else "null" | |
| done_val = str(done).lower() | |
| # Collapse newlines so the entire step fits on one line (spec requirement) | |
| action_clean = action.replace("\n", " ").replace("\r", "").strip() | |
| print( | |
| f"[STEP] step={step} action={action_clean} " | |
| f"reward={reward:.2f} done={done_val} error={err}", | |
| flush=True, | |
| ) | |
| def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None: | |
| rewards_str = ",".join(f"{r:.2f}" for r in rewards) if rewards else "" | |
| print( | |
| f"[END] success={str(success).lower()} steps={steps} " | |
| f"score={score:.3f} rewards={rewards_str}", | |
| flush=True, | |
| ) | |
| # --------------------------------------------------------------------------- | |
| # LLM / heuristic helpers | |
| # --------------------------------------------------------------------------- | |
| SYSTEM_PROMPT = textwrap.dedent( | |
| """ | |
| You are an expert SQL debugger. | |
| You will be shown a broken SQL query that contains typos or keyword errors. | |
| Fix ALL errors and return ONLY the corrected SQL query. | |
| No explanation, no markdown, no code blocks, no backticks. | |
| Common errors: FORM->FROM, WEHRE->WHERE, GRUP->GROUP, HAVNG->HAVING, | |
| ORDR->ORDER, INNE->INNER, LFT->LEFT, BETWEN->BETWEEN, DSC->DESC, SELCT->SELECT. | |
| """ | |
| ).strip() | |
| SQL_REPLACEMENTS = { | |
| "FORM": "FROM", | |
| "WEHRE": "WHERE", | |
| "WHER": "WHERE", | |
| "GRUP": "GROUP", | |
| "HAVNG": "HAVING", | |
| "ORDR": "ORDER", | |
| "INNE": "INNER", | |
| "LFT": "LEFT", | |
| "BETWEN": "BETWEEN", | |
| "DSC": "DESC", | |
| "SELCT": "SELECT", | |
| "LIMT": "LIMIT", | |
| "DPT_ID": "DEPT_ID", | |
| } | |
| def heuristic_correct_sql(query: str) -> str: | |
| """Deterministic fallback when the LLM is unavailable.""" | |
| corrected = query | |
| for broken, fixed in SQL_REPLACEMENTS.items(): | |
| corrected = re.sub( | |
| rf"\b{re.escape(broken)}\b", fixed, corrected, flags=re.IGNORECASE | |
| ) | |
| return corrected.strip() | |
| def get_model_action( | |
| client: Optional["OpenAI"], obs: dict, history: List[str] | |
| ) -> str: | |
| """Return a corrected SQL string. Falls back to heuristic on any failure.""" | |
| heuristic = heuristic_correct_sql(obs.get("broken_query", "")) | |
| if client is None: | |
| return heuristic | |
| history_block = "\n".join(history[-4:]) if history else "None" | |
| user_prompt = textwrap.dedent( | |
| f""" | |
| Broken SQL query: | |
| {obs.get("broken_query", "")} | |
| Schema context: {obs.get("schema_context") or "Not provided"} | |
| Error hint: {obs.get("error_hint") or "None"} | |
| Your previous attempt: {obs.get("previous_attempt") or "None"} | |
| Feedback: {obs.get("feedback") or "None"} | |
| Recent history: | |
| {history_block} | |
| Return ONLY the corrected SQL query. | |
| """ | |
| ).strip() | |
| try: | |
| completion = client.chat.completions.create( | |
| model=MODEL_NAME, | |
| messages=[ | |
| {"role": "system", "content": SYSTEM_PROMPT}, | |
| {"role": "user", "content": user_prompt}, | |
| ], | |
| temperature=0.2, | |
| max_tokens=300, | |
| stream=False, | |
| ) | |
| text = (completion.choices[0].message.content or "").strip() | |
| return text if text else heuristic | |
| except Exception as exc: | |
| print(f"[DEBUG] LLM call failed: {exc}", flush=True) | |
| return heuristic | |
| # --------------------------------------------------------------------------- | |
| # Episode runner | |
| # --------------------------------------------------------------------------- | |
| async def run_task(task_name: str) -> None: | |
| """ | |
| Run one full episode for `task_name`. | |
| The [END] log line is ALWAYS emitted via the finally block, even if an | |
| exception occurs mid-episode or the reset call fails. This is required | |
| by the hackathon spec to avoid disqualification. | |
| """ | |
| # Initialise all accumulators BEFORE the try so finally can always read them | |
| rewards: List[float] = [] | |
| history: List[str] = [] | |
| steps_taken = 0 | |
| score = 0.0 | |
| success = False | |
| # Build LLM client (best-effort; None means heuristic-only mode) | |
| client = None | |
| if OpenAI is not None and API_KEY not in {"", "dummy"}: | |
| try: | |
| client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY) | |
| except Exception as exc: | |
| print(f"[DEBUG] OpenAI client init failed: {exc}", flush=True) | |
| log_start(task_name, BENCHMARK, MODEL_NAME) | |
| http: Optional[httpx.AsyncClient] = None | |
| try: | |
| http = httpx.AsyncClient(base_url=ENV_URL, timeout=60.0) | |
| # --- reset --------------------------------------------------------- | |
| reset_failed = False | |
| obs: dict = {} | |
| try: | |
| reset_resp = await http.post("/reset", json={"difficulty": task_name}) | |
| reset_resp.raise_for_status() | |
| obs = reset_resp.json() | |
| except Exception as exc: | |
| print(f"[DEBUG] Reset failed: {exc}", flush=True) | |
| # Do NOT return here — fall through to finally so [END] is always logged | |
| reset_failed = True | |
| if not reset_failed: | |
| # --- step loop ------------------------------------------------- | |
| for step in range(1, MAX_STEPS + 1): | |
| # Get action (never raises — heuristic is the ultimate fallback) | |
| try: | |
| action_str = get_model_action(client, obs, history) | |
| except Exception as exc: | |
| print(f"[DEBUG] Model action failed: {exc}", flush=True) | |
| action_str = heuristic_correct_sql(obs.get("broken_query", "")) | |
| # Submit action to environment | |
| try: | |
| step_resp = await http.post( | |
| "/step", json={"corrected_query": action_str} | |
| ) | |
| step_resp.raise_for_status() | |
| result = step_resp.json() | |
| except Exception as exc: | |
| print(f"[DEBUG] Step {step} request failed: {exc}", flush=True) | |
| # Treat as a 0-reward terminal step so episode ends cleanly | |
| rewards.append(0.0) | |
| steps_taken = step | |
| log_step(step, action_str, 0.0, True, str(exc)) | |
| break | |
| obs = result.get("observation", obs) | |
| reward = float(result.get("reward", 0.0)) | |
| done = bool(result.get("done", False)) | |
| info = result.get("info") | |
| error = info.get("error") if isinstance(info, dict) else None | |
| rewards.append(reward) | |
| steps_taken = step | |
| history.append( | |
| f"Step {step}: attempt={action_str!r} reward={reward:+.2f}" | |
| ) | |
| log_step(step, action_str, reward, done, error) | |
| if done: | |
| break | |
| # Score = average reward across all steps, clamped to [0, 1] | |
| if rewards: | |
| score = min(max(sum(rewards) / len(rewards), 0.0), 1.0) | |
| success = score >= SUCCESS_SCORE_THRESHOLD | |
| except Exception as exc: | |
| # Catch-all for any unexpected error in the episode body | |
| print(f"[DEBUG] Unhandled episode error: {exc}", flush=True) | |
| finally: | |
| # Always close the HTTP client | |
| if http is not None: | |
| try: | |
| await http.aclose() | |
| except Exception as exc: | |
| print(f"[DEBUG] HTTP close error: {exc}", flush=True) | |
| # [END] MUST always be emitted — even after reset failure or exception | |
| log_end(success, steps_taken, score, rewards) | |
| # --------------------------------------------------------------------------- | |
| # Entry point | |
| # --------------------------------------------------------------------------- | |
| async def main() -> None: | |
| """ | |
| Run tasks according to SQL_ENV_TASK. | |
| If SQL_ENV_TASK is a single valid difficulty, run only that task. | |
| Otherwise run all three in sequence so all 3 tasks produce scores. | |
| """ | |
| try: | |
| difficulties = ( | |
| (TASK_NAME,) | |
| if TASK_NAME in {"easy", "medium", "hard"} | |
| else ("easy", "medium", "hard") | |
| ) | |
| for difficulty in difficulties: | |
| await run_task(difficulty) | |
| print("", flush=True) # blank line separator between tasks | |
| except Exception as exc: | |
| print(f"[DEBUG] Main loop error: {exc}", flush=True) | |
| if __name__ == "__main__": | |
| try: | |
| asyncio.run(main()) | |
| except KeyboardInterrupt: | |
| pass | |
| except Exception as exc: | |
| print(f"[DEBUG] Fatal error: {exc}", flush=True) | |
| finally: | |
| sys.exit(0) | |