Spaces:
Sleeping
Sleeping
File size: 2,551 Bytes
eb70591 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | import os
import asyncio
from openai import AsyncOpenAI
from client import SQLDebugEnv, SQLFixAction
# Environment variables
API_BASE_URL = os.environ.get("API_BASE_URL", "https://api.openai.com/v1")
MODEL_NAME = os.environ.get("MODEL_NAME", "gpt-4o-mini")
HF_TOKEN = os.environ.get("HF_TOKEN", "")
ENV_BASE_URL = os.environ.get("ENV_BASE_URL", "http://localhost:8000")
async def run_inference():
"""
Baseline inference script.
Connects to the environment, resets an episode, and uses an LLM to generate a fix.
"""
client = AsyncOpenAI(
base_url=API_BASE_URL,
api_key=HF_TOKEN or "dummy-key"
)
async with SQLDebugEnv(base_url=ENV_BASE_URL) as env:
print("Starting new episode...")
result = await env.reset(task_group="easy_syntax_fix")
obs = result.observation
print(f"Task: {obs.task_description}")
print(f"Buggy Query: {obs.buggy_query}")
print(f"Error: {obs.error_message}")
prompt = f"""
You are an expert SQL engineer. Your task is to fix the following buggy query.
Task Description: {obs.task_description}
Database Schema:
{obs.schema_sql}
Sample Data:
{obs.sample_data}
Buggy Query: {obs.buggy_query}
Error Message: {obs.error_message}
Return ONLY the corrected SQL query. Do not include markdown formatting or explanations.
"""
try:
response = await client.chat.completions.create(
model=MODEL_NAME,
messages=[{"role": "user", "content": prompt}],
temperature=0.0
)
corrected_query = response.choices[0].message.content.strip()
if corrected_query.startswith("```sql"):
corrected_query = corrected_query[6:]
if corrected_query.endswith("```"):
corrected_query = corrected_query[:-3]
corrected_query = corrected_query.strip()
step_resp = await env.step(SQLFixAction(query=corrected_query))
print(f"Reward: {step_resp.reward}")
print(f"Success: {step_resp.success}")
print(f"Feedback: {step_resp.observation.last_attempt_result}")
except Exception as e:
print(f"LLM API Error: {e}")
if __name__ == "__main__":
asyncio.run(run_inference())
|