Spaces:
Sleeping
Sleeping
| import asyncio | |
| import os | |
| from typing import List, Optional | |
| from openai import OpenAI | |
| from client import SQLDebuggerClient | |
| from models import SQLAction | |
| # Step 1: Constants & env vars | |
| HF_TOKEN = os.getenv("HF_TOKEN") | |
| API_KEY = HF_TOKEN or os.getenv("API_KEY") | |
| API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1") | |
| MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct") | |
| LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME") | |
| ENV_URL = os.getenv("ENV_URL", "http://localhost:8000") | |
| BENCHMARK = "sql_debugger" | |
| MAX_STEPS = 3 # 3 attempts per task, matches our environment | |
| TEMPERATURE = 0.3 # low temperature = more deterministic SQL output | |
| MAX_TOKENS = 256 # SQL queries are short, don't need much | |
| # Step 2: Log functions — exact format required by hackathon grading | |
| 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: | |
| done_val = str(done).lower() # Python's True/False → "true"/"false" | |
| error_val = error if error else "null" # None → "null" | |
| 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: | |
| 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) | |
| # Step 3: Prompts | |
| SYSTEM_PROMPT = """You are a SQL debugging expert. You will be given: | |
| - A broken SQL query that has errors | |
| - The database schema (CREATE TABLE statements) | |
| - A description of what's wrong | |
| - The expected output the query should produce | |
| Your job is to fix the SQL query so it produces the expected output. | |
| Rules: | |
| - Return ONLY the corrected SQL query, nothing else | |
| - No explanations, no markdown, no code blocks | |
| - Just the raw SQL query ending with a semicolon""" | |
| def build_user_prompt(observation, feedback_history: List[str]) -> str: | |
| """Build the prompt we send to the LLM for each attempt.""" | |
| schema_str = "\n".join(observation.table_schema) | |
| prompt = ( | |
| f"Broken SQL query:\n{observation.broken_query}\n\n" | |
| f"Database schema:\n{schema_str}\n\n" | |
| f"Description: {observation.description}\n\n" | |
| f"Expected output: {observation.expected_output}" | |
| ) | |
| # If the agent has tried before, include feedback so it can learn | |
| if feedback_history: | |
| prompt += "\n\nYour previous attempts:" | |
| for entry in feedback_history: | |
| prompt += f"\n{entry}" | |
| prompt += "\n\nFix the query based on the feedback above." | |
| return prompt | |
| # Step 4: LLM call function | |
| def get_corrected_sql(client: OpenAI, observation, feedback_history: List[str]) -> str: | |
| """Send the puzzle to the LLM, get back a corrected SQL query.""" | |
| user_prompt = build_user_prompt(observation, feedback_history) | |
| try: | |
| completion = client.chat.completions.create( | |
| model=MODEL_NAME, | |
| messages=[ | |
| {"role": "system", "content": SYSTEM_PROMPT}, | |
| {"role": "user", "content": user_prompt}, | |
| ], | |
| temperature=TEMPERATURE, | |
| max_tokens=MAX_TOKENS, | |
| stream=False, | |
| ) | |
| sql = (completion.choices[0].message.content or "").strip() | |
| # Clean up in case the LLM wraps it in markdown code blocks | |
| if sql.startswith("```"): | |
| sql = sql.split("\n", 1)[-1] # remove first line (```sql) | |
| sql = sql.rsplit("```", 1)[0] # remove closing ``` | |
| sql = sql.strip() | |
| return sql if sql else "SELECT 1;" | |
| except Exception as e: | |
| print(f"[DEBUG] LLM request failed: {e}", flush=True) | |
| return "SELECT 1;" | |
| # Step 5: Main loop | |
| async def main() -> None: | |
| # Create the LLM client (sync) and environment client (async) | |
| llm_client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY) | |
| if LOCAL_IMAGE_NAME: | |
| env = await SQLDebuggerClient.from_docker_image(LOCAL_IMAGE_NAME) | |
| else: | |
| env = SQLDebuggerClient(base_url=ENV_URL) | |
| await env.connect() | |
| num_tasks = 9 # we have 9 tasks in our bank | |
| try: | |
| for task_num in range(num_tasks): | |
| # Start a new task — get the puzzle | |
| result = await env.reset() | |
| obs = result.observation | |
| task_id = f"task_{task_num + 1}" | |
| log_start(task=task_id, env=BENCHMARK, model=MODEL_NAME) | |
| rewards: List[float] = [] | |
| feedback_history: List[str] = [] | |
| steps_taken = 0 | |
| final_score = 0.0 | |
| # Up to 3 attempts per task | |
| for step in range(1, MAX_STEPS + 1): | |
| # Ask the LLM to fix the broken SQL | |
| corrected_sql = get_corrected_sql(llm_client, obs, feedback_history) | |
| # Send the fix to the environment, get score | |
| result = await env.step(SQLAction(corrected_query=corrected_sql)) | |
| obs = result.observation | |
| reward = result.reward or 0.0 | |
| done = result.done | |
| rewards.append(reward) | |
| steps_taken = step | |
| final_score = max(final_score, reward) # best attempt counts | |
| log_step( | |
| step=step, | |
| action=corrected_sql, | |
| reward=reward, | |
| done=done, | |
| error=obs.error_message, | |
| ) | |
| # Build feedback for next attempt (if any) | |
| feedback_history.append( | |
| f"Attempt {step}: '{corrected_sql}' → score={reward:.2f}, error={obs.error_message or 'none'}" | |
| ) | |
| if done: | |
| break | |
| success = final_score >= 0.99 | |
| log_end(success=success, steps=steps_taken, score=final_score, rewards=rewards) | |
| finally: | |
| try: | |
| await env.close() | |
| except Exception as e: | |
| print(f"[DEBUG] env.close() error: {e}", flush=True) | |
| if __name__ == "__main__": | |
| asyncio.run(main()) | |