inference.py

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())