Spaces:
Sleeping
Sleeping
| 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()) | |