import argparse import json import os import sys import textwrap from typing import List, Optional from dotenv import load_dotenv load_dotenv() from openai import OpenAI from client import SqlSandboxEnv from models import SqlSandboxAction # --------------------------------------------------------------------------- # Ensure required env vars have fallbacks so OpenAI client never gets None # --------------------------------------------------------------------------- API_BASE_URL = os.getenv("API_BASE_URL") or "https://router.huggingface.co/v1" API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY") or os.getenv("OPENAI_API_KEY") MODEL_NAME = os.getenv("MODEL_NAME") or "gpt-4o-mini" BENCHMARK = "sql_sandbox" SYSTEM_PROMPT = textwrap.dedent(""" You are a data engineering assistant working inside a SQLite sandbox. You can execute two types of actions: 1. {"tool": "sql", "command": ""} 2. {"tool": "python", "command": ""} Rules: 1 Respond with EXACTLY ONE JSON object per turn no markdown, no explanation. 2 In Python code, the variables `conn` (sqlite3.Connection) and `cursor` (sqlite3.Cursor) are already available. Do NOT call sqlite3.connect(). 3 SQLite STRFTIME months are zero-padded: use '01' not '1', or use LIKE '2024-01-%'. 4 When you believe the task is fully complete, send: {"tool": "sql", "command": "SELECT 'DONE'"} """).strip() 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: error_val = error.replace("\n", " ") if error else "null" done_val = str(done).lower() 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:.2f} rewards={rewards_str}", flush=True) def _run_task_agent(client_llm: OpenAI, base_url: str, task_id: str, max_turns: int = 15) -> float: rewards: List[float] = [] step_count = 0 final_score = 0.0 # Fallback response for API failures fallback_action = '{"tool": "sql", "command": "SELECT \'DONE\'"}' with SqlSandboxEnv(base_url=base_url).sync() as env: try: reset_resp = env.reset(task_id=task_id) task_desc = reset_resp.observation.task_description except Exception as e: print(f"[DEBUG] env.reset() error for task {task_id}: {e}", flush=True) return 0.0 messages = [ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": f"Task: {task_desc}\n\nBegin."}, ] log_start(task=task_id, env=BENCHMARK, model=MODEL_NAME) for turn in range(1, max_turns + 1): # 1. Ask the LLM, wrapped in try...except try: 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() except Exception as exc: print(f"[DEBUG] Model request failed: {exc}", flush=True) assistant_msg = fallback_action # 2. Parse action JSON 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 try: step_resp = env.step(SqlSandboxAction(tool=tool, command=command)) except Exception as exc: print(f"[DEBUG] env.step() error: {exc}", flush=True) break reward = step_resp.reward or 0.0 done = step_resp.done output = step_resp.observation.output or "" error = step_resp.observation.error or "" rewards.append(reward) step_count += 1 action_str = json.dumps({"tool": tool, "command": command}) log_step(step=step_count, action=action_str, reward=reward, done=done, error=error) 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}) raw_score = sum(rewards) final_score = max(0.01, min(0.99, float(raw_score))) success = final_score >= 0.99 log_end(success=success, steps=step_count, score=final_score, rewards=rewards) return final_score 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", ) parser.add_argument( "--max-turns", type=int, default=15, help="Maximum agent turns per task (default: 15)", ) args = parser.parse_args() if not API_KEY: print("ERROR: HF_TOKEN (or OPENAI_API_KEY) environment variable is not set.", flush=True) client_llm = OpenAI( api_key=API_KEY or "dummy_key", base_url=API_BASE_URL, ) for task in [f"task{i}" for i in range(1, 7)]: _run_task_agent(client_llm, args.url, task, args.max_turns) if __name__ == "__main__": main()