Upload folder using huggingface_hub

baseline.py

@@ -1,32 +1,48 @@
+"""
+Baseline inference script for the SQL Query Optimizer OpenEnv.
+
+Uses the OpenAI API client to run a model against the environment
+and produce reproducible baseline scores on all 3 tasks.
+
+Usage:
+    export OPENAI_API_KEY=sk-...
+    python baseline.py
+"""
+
 import os
 from openai import OpenAI
 from env.environment import SQLEnv
 from env.models import Action
 
+
 def run_task(env: SQLEnv, task_id: int) -> float:
     client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
-    obs = env.reset(task_id)
-    
+    obs = env.reset(task_id=task_id)
+
     messages = [
-        {"role": "system", "content": "You are an expert SQL DBA. You rewrite SQL queries to be correct, optimized, and performant."}
+        {
+            "role": "system",
+            "content": (
+                "You are an expert SQL DBA. You rewrite SQL queries "
+                "to be correct, optimized, and performant."
+            ),
+        }
     ]
-    
-    prompt = f"""
-Task # {obs.task_id}
+
+    prompt = f"""Task #{obs.task_id}
 Original Query: {obs.query}
 Database Schema Context: {obs.schema_context}
 Hint: {obs.hint}
 
-Please provide the optimized query. Output ONLY the raw SQL query, no markdown formatting, no explanation.
-"""
-    
+Please provide the optimized query. Output ONLY the raw SQL query, no markdown formatting, no explanation."""
+
     messages.append({"role": "user", "content": prompt.strip()})
-    
+
     try:
         response = client.chat.completions.create(
             model="gpt-3.5-turbo",
             messages=messages,
-            temperature=0.0
+            temperature=0.0,
         )
         rewritten_query = response.choices[0].message.content.strip()
         if rewritten_query.startswith("```sql"):
@@ -37,30 +53,32 @@ Please provide the optimized query. Output ONLY the raw SQL query, no markdown f
     except Exception as e:
         print(f"Error calling OpenAI API: {e}")
         rewritten_query = obs.query
-        
+
     action = Action(
         rewritten_query=rewritten_query,
         explanation="Baseline inference using LLM",
-        is_done=True
+        is_done=True,
     )
-    
-    _, reward, done, info = env.step(action)
-    return env.final_grader_score
+
+    result_obs = env.step(action)
+    return result_obs.reward
+
 
 def run_all_tasks():
     if not os.environ.get("OPENAI_API_KEY"):
         raise ValueError("OPENAI_API_KEY environment variable is required.")
-        
+
     env = SQLEnv()
     scores = {}
     for task_id in [1, 2, 3]:
         print(f"Running baseline for Task {task_id}...")
         score = run_task(env, task_id)
         scores[task_id] = score
-        print(f"Task {task_id} Grader Score: {score}")
-        
+        print(f"Task {task_id} Score: {score}")
+
     return scores
 
+
 if __name__ == "__main__":
     try:
         scores = run_all_tasks()

env/environment.py

@@ -1,9 +1,19 @@
-from typing import Tuple, Dict, Any, List
+from typing import Optional, Dict, Any
+from uuid import uuid4
+
+from openenv.core.env_server.interfaces import Environment
+from openenv.core.env_server.types import State
+
 from .models import Observation, Action, Reward
 from .tasks import TASKS, grade_action, get_task
 from .reward import compute_reward
 
-class SQLEnv:
+
+class SQLEnv(Environment):
+    """SQL Query Optimizer Environment following the OpenEnv interface."""
+
+    SUPPORTS_CONCURRENT_SESSIONS: bool = True
+
     def __init__(self):
         self.current_task_id = None
         self.task = None
@@ -13,12 +23,19 @@ class SQLEnv:
         self.cumulative_score = 0.0
         self.previous_grader_score = 0.0
         self.final_grader_score = 0.0
-        
-    def reset(self, task_id: int) -> Observation:
+        self._state = State(episode_id=str(uuid4()), step_count=0)
+
+    def reset(
+        self,
+        seed: Optional[int] = None,
+        episode_id: Optional[str] = None,
+        task_id: int = 1,
+        **kwargs: Any,
+    ) -> Observation:
         task = get_task(task_id)
         if not task:
             raise ValueError(f"Task {task_id} not found.")
-            
+
         self.current_task_id = task_id
         self.task = task
         self.step_number = 1
@@ -27,80 +44,88 @@ class SQLEnv:
         self.cumulative_score = 0.0
         self.previous_grader_score = 0.0
         self.final_grader_score = 0.0
-        
+        self._state = State(
+            episode_id=episode_id or str(uuid4()),
+            step_count=0,
+        )
+
         obs = Observation(
             task_id=self.current_task_id,
             query=self.task["initial_query"],
             schema_context=self.task["schema_context"],
             hint=self.task["hint"],
             step_number=self.step_number,
-            max_steps=self.max_steps
+            max_steps=self.max_steps,
+            reward=0.0,
+            done=False,
         )
         self.history.append({"step": 0, "type": "reset", "observation": obs.model_dump()})
         return obs
-        
-    def step(self, action: Action) -> Tuple[Observation, Reward, bool, Dict[str, Any]]:
+
+    def step(
+        self,
+        action: Action,
+        timeout_s: Optional[float] = None,
+        **kwargs: Any,
+    ) -> Observation:
         if not self.task:
             raise RuntimeError("Environment not initialized. Call reset() first.")
-            
-        grader_score, breakdown, feedback = grade_action(self.current_task_id, action.rewritten_query)
+
+        grader_score, breakdown, feedback = grade_action(
+            self.current_task_id, action.rewritten_query
+        )
         action_valid = len(action.rewritten_query.strip()) > 0
-        
+
         done = action.is_done or self.step_number >= self.max_steps
-        
+
         step_reward = compute_reward(
             grader_score=grader_score,
             previous_score=self.previous_grader_score,
             step_number=self.step_number,
             max_steps=self.max_steps,
             is_done=done,
-            action_valid=action_valid
+            action_valid=action_valid,
         )
-        
+
         self.cumulative_score += step_reward
         self.previous_grader_score = grader_score
-        
-        reward = Reward(
-            score=step_reward,
-            breakdown=breakdown,
-            feedback=feedback
-        )
-        
+
+        info = {
+            "cumulative_score": self.cumulative_score,
+            "grader_score": grader_score,
+            "breakdown": breakdown,
+            "feedback": feedback,
+        }
+
+        if done:
+            self.final_grader_score = grader_score
+
+        self._state.step_count += 1
+
         obs = Observation(
             task_id=self.current_task_id,
             query=action.rewritten_query,
             schema_context=self.task["schema_context"],
-            hint=self.task["hint"], 
+            hint=self.task["hint"],
             step_number=self.step_number + 1,
-            max_steps=self.max_steps
+            max_steps=self.max_steps,
+            reward=step_reward,
+            done=done,
+            metadata=info,
         )
-        
-        info = {
-            "cumulative_score": self.cumulative_score,
-            "grader_score": grader_score
-        }
-        
-        if done:
-            self.final_grader_score = grader_score
-            
+
         self.history.append({
             "step": self.step_number,
             "type": "step",
             "action": action.model_dump(),
-            "reward": reward.model_dump(),
+            "reward": step_reward,
             "done": done,
-            "info": info
+            "info": info,
         })
-        
+
         self.step_number += 1
-        return obs, reward, done, info
-        
-    def state(self) -> Dict[str, Any]:
-        return {
-            "current_task_id": self.current_task_id,
-            "step_number": self.step_number,
-            "max_steps": self.max_steps,
-            "cumulative_score": self.cumulative_score,
-            "final_grader_score": self.final_grader_score,
-            "history": self.history
-        }
+        return obs
+
+    @property
+    def state(self) -> State:
+        return self._state

env/models.py

@@ -1,20 +1,28 @@
-from typing import Optional, Dict
-from pydantic import BaseModel, Field
+from typing import Optional, Dict, Any
+from pydantic import Field
+from openenv.core.env_server.types import Action as BaseAction, Observation as BaseObservation
 
-class Observation(BaseModel):
-    task_id: int = Field(description="The ID of the task to perform.")
-    query: str = Field(description="The SQL query to review and optimize.")
-    schema_context: str = Field(description="The database schema context for the query, such as CREATE TABLE statements.")
-    hint: Optional[str] = Field(default=None, description="An optional natural-language hint or description of the problem.")
-    step_number: int = Field(description="The current step number in the episode (1-indexed).")
-    max_steps: int = Field(description="The maximum allowed steps for this task.")
 
-class Action(BaseModel):
-    rewritten_query: str = Field(description="The rewritten, optimized SQL query.")
-    explanation: str = Field(description="A brief explanation of the changes made and why they improve the query.")
-    is_done: bool = Field(description="Set to true if you are finished and want to submit the query for final scoring.")
+class Observation(BaseObservation):
+    task_id: int = Field(default=0, description="The ID of the task to perform.")
+    query: str = Field(default="", description="The SQL query to review and optimize.")
+    schema_context: str = Field(default="", description="The database schema context.")
+    hint: Optional[str] = Field(default=None, description="An optional natural-language hint.")
+    step_number: int = Field(default=0, description="The current step number in the episode.")
+    max_steps: int = Field(default=0, description="The maximum allowed steps for this task.")
 
-class Reward(BaseModel):
-    score: float = Field(description="The overall score for the episode (0.0 to 1.0).")
-    breakdown: Dict[str, float] = Field(default_factory=dict, description="A breakdown of the score by sub-criteria.")
-    feedback: str = Field(description="Specific feedback on the rewritten query or action taken.")
+
+class Action(BaseAction):
+    rewritten_query: str = Field(default="", description="The rewritten, optimized SQL query.")
+    explanation: str = Field(default="", description="A brief explanation of the changes.")
+    is_done: bool = Field(default=False, description="Set to true to submit for final scoring.")
+
+
+class Reward:
+    def __init__(self, score: float = 0.0, breakdown: Dict[str, float] = None, feedback: str = ""):
+        self.score = score
+        self.breakdown = breakdown or {}
+        self.feedback = feedback
+
+    def model_dump(self):
+        return {"score": self.score, "breakdown": self.breakdown, "feedback": self.feedback}

inference.py

@@ -0,0 +1,89 @@
+"""
+Baseline inference script for the SQL Query Optimizer OpenEnv.
+
+Uses the OpenAI API client to run a model against the environment
+and produce reproducible baseline scores on all 3 tasks.
+
+Usage:
+    export OPENAI_API_KEY=sk-...
+    python inference.py
+"""
+
+import os
+from openai import OpenAI
+from env.environment import SQLEnv
+from env.models import Action
+
+
+def run_task(env: SQLEnv, task_id: int) -> float:
+    client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
+    obs = env.reset(task_id=task_id)
+
+    messages = [
+        {
+            "role": "system",
+            "content": (
+                "You are an expert SQL DBA. You rewrite SQL queries "
+                "to be correct, optimized, and performant."
+            ),
+        }
+    ]
+
+    prompt = f"""Task #{obs.task_id}
+Original Query: {obs.query}
+Database Schema Context: {obs.schema_context}
+Hint: {obs.hint}
+
+Please provide the optimized query. Output ONLY the raw SQL query, no markdown formatting, no explanation."""
+
+    messages.append({"role": "user", "content": prompt.strip()})
+
+    try:
+        response = client.chat.completions.create(
+            model="gpt-3.5-turbo",
+            messages=messages,
+            temperature=0.0,
+        )
+        rewritten_query = response.choices[0].message.content.strip()
+        if rewritten_query.startswith("```sql"):
+            rewritten_query = rewritten_query[6:]
+        if rewritten_query.endswith("```"):
+            rewritten_query = rewritten_query[:-3]
+        rewritten_query = rewritten_query.strip()
+    except Exception as e:
+        print(f"Error calling OpenAI API: {e}")
+        rewritten_query = obs.query
+
+    action = Action(
+        rewritten_query=rewritten_query,
+        explanation="Baseline inference using LLM",
+        is_done=True,
+    )
+
+    result_obs = env.step(action)
+    return result_obs.reward
+
+
+def run_all_tasks():
+    if not os.environ.get("OPENAI_API_KEY"):
+        raise ValueError("OPENAI_API_KEY environment variable is required.")
+
+    env = SQLEnv()
+    scores = {}
+    for task_id in [1, 2, 3]:
+        print(f"Running baseline for Task {task_id}...")
+        score = run_task(env, task_id)
+        scores[task_id] = score
+        print(f"Task {task_id} Score: {score}")
+
+    return scores
+
+
+if __name__ == "__main__":
+    try:
+        scores = run_all_tasks()
+        print("\nBaseline Evaluation Results:")
+        for t, s in scores.items():
+            print(f"Task {t}: {s}/1.0")
+    except Exception as e:
+        print(f"Baseline Evaluation Failed: {e}")

models.py

@@ -1,20 +1,28 @@
-from typing import Optional, Dict
-from pydantic import BaseModel, Field
+from typing import Optional, Dict, Any
+from pydantic import Field
+from openenv.core.env_server.types import Action as BaseAction, Observation as BaseObservation
 
-class Observation(BaseModel):
-    task_id: int = Field(description="The ID of the task to perform.")
-    query: str = Field(description="The SQL query to review and optimize.")
-    schema_context: str = Field(description="The database schema context for the query, such as CREATE TABLE statements.")
-    hint: Optional[str] = Field(default=None, description="An optional natural-language hint or description of the problem.")
-    step_number: int = Field(description="The current step number in the episode (1-indexed).")
-    max_steps: int = Field(description="The maximum allowed steps for this task.")
 
-class Action(BaseModel):
-    rewritten_query: str = Field(description="The rewritten, optimized SQL query.")
-    explanation: str = Field(description="A brief explanation of the changes made and why they improve the query.")
-    is_done: bool = Field(description="Set to true if you are finished and want to submit the query for final scoring.")
+class Observation(BaseObservation):
+    task_id: int = Field(default=0, description="The ID of the task to perform.")
+    query: str = Field(default="", description="The SQL query to review and optimize.")
+    schema_context: str = Field(default="", description="The database schema context.")
+    hint: Optional[str] = Field(default=None, description="An optional natural-language hint.")
+    step_number: int = Field(default=0, description="The current step number in the episode.")
+    max_steps: int = Field(default=0, description="The maximum allowed steps for this task.")
 
-class Reward(BaseModel):
-    score: float = Field(description="The overall score for the episode (0.0 to 1.0).")
-    breakdown: Dict[str, float] = Field(default_factory=dict, description="A breakdown of the score by sub-criteria.")
-    feedback: str = Field(description="Specific feedback on the rewritten query or action taken.")
+
+class Action(BaseAction):
+    rewritten_query: str = Field(default="", description="The rewritten, optimized SQL query.")
+    explanation: str = Field(default="", description="A brief explanation of the changes.")
+    is_done: bool = Field(default=False, description="Set to true to submit for final scoring.")
+
+
+class Reward:
+    def __init__(self, score: float = 0.0, breakdown: Dict[str, float] = None, feedback: str = ""):
+        self.score = score
+        self.breakdown = breakdown or {}
+        self.feedback = feedback
+
+    def model_dump(self):
+        return {"score": self.score, "breakdown": self.breakdown, "feedback": self.feedback}

server/app.py

@@ -12,26 +12,39 @@ app = create_app(
     env=SQLEnv,
     action_cls=Action,
     observation_cls=Observation,
-    env_name="sql-query-optimizer"
+    env_name="sql-query-optimizer",
 )
 
+
 @app.get("/tasks")
 async def get_tasks():
     action_schema = Action.model_json_schema()
     task_list = [{"id": k, **v} for k, v in TASKS.items()]
     return {
         "tasks": task_list,
-        "action_schema": action_schema
+        "action_schema": action_schema,
     }
 
+
 class BaselineResponse(BaseModel):
     scores: Dict[int, float]
 
+
 @app.post("/baseline", response_model=BaselineResponse)
 async def run_baseline():
     import baseline
+
     try:
         scores = baseline.run_all_tasks()
         return BaselineResponse(scores=scores)
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
+
+
+def main(host: str = "0.0.0.0", port: int = 7860):
+    import uvicorn
+    uvicorn.run(app, host=host, port=port)
+
+
+if __name__ == "__main__":
+    main()