inference.py

"""
Baseline Inference Script for SQL/Data Cleaning Sandbox  OpenAI Edition.

Uses OpenAI (gpt-4o) to solve all three tasks and prints reproducible
scores via the OpenEnv WebSocket client.

Usage:
    set HF_TOKEN=sk-...          # Windows
    export HF_TOKEN=sk-...       # Linux/macOS
    python inference.py                    # local server
    python inference.py --url https://...  # remote server
"""

import argparse
import json
import os
import sys

from dotenv import load_dotenv
load_dotenv()

from openai import OpenAI

from client import SqlSandboxEnv
from models import SqlSandboxAction


# ---------------------------------------------------------------------------
# System prompt shared across all tasks
# ---------------------------------------------------------------------------
SYSTEM_PROMPT = """\
You are a data engineering assistant working inside a SQLite sandbox.

You can execute two types of actions:
1. {"tool": "sql",    "command": "<SQL query>"}
2. {"tool": "python", "command": "<Python code>"}

Rules:
- Respond with EXACTLY ONE JSON object per turn  no markdown, no explanation.
- In Python code, the variables `conn` (sqlite3.Connection) and `cursor`
  (sqlite3.Cursor) are already available. Do NOT call sqlite3.connect().
- SQLite STRFTIME months are zero-padded: use '01' not '1', or use LIKE '2024-01-%'.
- When you believe the task is fully complete, send:
  {"tool": "sql", "command": "SELECT 'DONE'"}
"""


# ---------------------------------------------------------------------------
# Core agent loop  one task, one WebSocket session
# ---------------------------------------------------------------------------
def _run_task_agent(base_url: str, task_id: str, max_turns: int = 15) -> float:
    """
    Open a fresh WebSocket session, reset the environment to the given task,
    then run an LLM agent loop until done or max_turns is reached.
    Returns the final reward (0.0  1.0).
    """
    api_key = os.environ.get("HF_TOKEN") or os.environ.get("OPENAI_API_KEY")
    api_base_url = os.environ.get("API_BASE_URL")
    model_name = os.environ.get("MODEL_NAME", "gpt-4o")

    client_llm = OpenAI(
        api_key=api_key,
        base_url=api_base_url,
    )
    final_reward = 0.0

    # Each task gets its own WebSocket session to avoid state leakage
    with SqlSandboxEnv(base_url=base_url).sync() as env:
        # reset() with task_id seeds the correct DB table for this task
        reset_resp = env.reset(task_id=task_id)
        task_desc = reset_resp.observation.task_description

        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user",   "content": f"Task: {task_desc}\n\nBegin."},
        ]

        print(f"\n  --- Session: {task_id} ---")

        for turn in range(max_turns):
            # 1. Ask the LLM
            response = client_llm.chat.completions.create(
                model=model_name,
                messages=messages,
                temperature=0.0,
                max_tokens=512,
            )
            assistant_msg = response.choices[0].message.content.strip()

            # 2. Parse action JSON (handle optional markdown fences)
            try:
                raw = assistant_msg
                if raw.startswith("```"):
                    raw = raw.split("```")[1]
                    if raw.startswith("json"):
                        raw = raw[4:]
                action_data = json.loads(raw)
                tool    = action_data["tool"]
                command = action_data["command"]
            except (json.JSONDecodeError, KeyError):
                # Feed parse error back to LLM, do NOT count as a step
                messages.append({"role": "assistant", "content": assistant_msg})
                messages.append({
                    "role": "user",
                    "content": (
                        'Invalid JSON. Reply with exactly one JSON object:\n'
                        '{"tool": "sql" | "python", "command": "..."}'
                    ),
                })
                continue

            # 3. Execute the action via OpenEnv step()
            step_resp = env.step(SqlSandboxAction(tool=tool, command=command))

            reward = step_resp.reward or 0.0
            done   = step_resp.done
            output = step_resp.observation.output or ""
            error  = step_resp.observation.error  or ""

            final_reward = reward
            print(f"  [Turn {turn+1:02d}] tool={tool:<6} | reward={reward:.4f} | done={done}")

            if done:
                break

            # 4. Feed result back to LLM for the next turn
            messages.append({"role": "assistant", "content": assistant_msg})
            feedback = f"Output:\n{output[:1500]}"
            if error:
                feedback += f"\nError:\n{error[:500]}"
            feedback += f"\nReward so far: {reward:.4f}"
            messages.append({"role": "user", "content": feedback})

    return final_reward


# ---------------------------------------------------------------------------
# Per-difficulty entry points (called by main, importable for custom use)
# ---------------------------------------------------------------------------
def easy_run(base_url: str, max_turns: int = 15) -> float:
    print(f"\n{'='*50}\nRunning task: easy\n{'='*50}")
    score = _run_task_agent(base_url, "easy", max_turns)
    print(f"  Final score: {score:.4f}")
    return score


def med_run(base_url: str, max_turns: int = 15) -> float:
    print(f"\n{'='*50}\nRunning task: medium\n{'='*50}")
    score = _run_task_agent(base_url, "medium", max_turns)
    print(f"  Final score: {score:.4f}")
    return score


def hard_run(base_url: str, max_turns: int = 15) -> float:
    print(f"\n{'='*50}\nRunning task: hard\n{'='*50}")
    score = _run_task_agent(base_url, "hard", max_turns)
    print(f"  Final score: {score:.4f}")
    return score


# ---------------------------------------------------------------------------
# CLI entry point
# ---------------------------------------------------------------------------
def main():
    parser = argparse.ArgumentParser(
        description="OpenAI baseline inference for the SQL/Data Cleaning Sandbox"
    )
    parser.add_argument(
        "--url",
        default="http://localhost:7860",
        help="Base URL of the running environment server (default: http://localhost:7860)",
    )
    parser.add_argument(
        "--max-turns",
        type=int,
        default=15,
        help="Maximum agent turns per task (default: 15)",
    )
    args = parser.parse_args()

    if not os.environ.get("HF_TOKEN") and not os.environ.get("OPENAI_API_KEY"):
        print("ERROR: HF_TOKEN (or OPENAI_API_KEY) environment variable is not set per checklist.")
        sys.exit(1)

    results: dict[str, float] = {}
    results["easy"]   = easy_run(args.url, args.max_turns)
    results["medium"] = med_run(args.url,  args.max_turns)
    results["hard"]   = hard_run(args.url, args.max_turns)

    avg = sum(results.values()) / len(results)
    print(f"\n{'='*50}")
    print("RESULTS SUMMARY")
    print(f"{'='*50}")
    for task_id, score in results.items():
        print(f"  {task_id:<10}: {score:.4f}")
    print(f"  {'average':<10}: {avg:.4f}")


if __name__ == "__main__":
    main()