feat: implement SQL agent environment with OpenEnv integration, SQLite management, and automated task grading

Dockerfile

@@ -0,0 +1,25 @@
+FROM python:3.10-slim
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y curl && rm -rf /var/lib/apt/lists/*
+
+WORKDIR /app
+
+# Install dependencies
+COPY requirements.txt /tmp/requirements.txt
+RUN pip install --no-cache-dir -r /tmp/requirements.txt && rm /tmp/requirements.txt
+
+# Copy environment code
+COPY . /app/
+
+# Environment configurations setup
+ENV PYTHONPATH="/app"
+
+# Health check
+HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \
+    CMD curl -f http://localhost:8000/docs || exit 1
+
+EXPOSE 8000
+
+# Run server with updated root OOP Path
+CMD ["uvicorn", "server.app:app", "--host", "0.0.0.0", "--port", "8000"]

__init__.py

@@ -0,0 +1,4 @@
+from models import SqlAction, SqlObservation, SqlState
+from client import SqlEnvClient
+
+__all__ = ["SqlAction", "SqlObservation", "SqlState", "SqlEnvClient"]

client.py

@@ -0,0 +1,18 @@
+from openenv.core import EnvClient
+from openenv.core.client_types import StepResult
+from models import SqlAction, SqlObservation, SqlState
+
+class SqlEnvClient(EnvClient[SqlAction, SqlObservation, SqlState]):
+    def _step_payload(self, action: SqlAction) -> dict:
+        return {"query": action.query}
+
+    def _parse_result(self, payload: dict) -> StepResult[SqlObservation]:
+        obs = SqlObservation(**payload["observation"])
+        return StepResult(
+            observation=obs,
+            reward=payload.get("reward"),
+            done=payload.get("done", False),
+        )
+
+    def _parse_state(self, payload: dict) -> SqlState:
+        return SqlState(**payload)

core/__init__.py

@@ -0,0 +1,5 @@
+# Core components exports
+from core.grader import DataFrameGrader
+from core.tasks import SqlTask, EnvironmentTasks
+
+__all__ = ["DataFrameGrader", "SqlTask", "EnvironmentTasks"]

core/grader.py

@@ -0,0 +1,38 @@
+import pandas as pd
+
+class DataFrameGrader:
+    """Handles logic for comparing agent SQL result DataFrames against expected DataFrames."""
+    
+    @staticmethod
+    def grade(agent_df: pd.DataFrame, expected_df: pd.DataFrame) -> float:
+        """
+        Grades a DataFrame result. Returns 0.0 to 1.0.
+        Checks for exact match (columns, rows, data types).
+        Falls back to unordered match if the structure is the same but sorting is off.
+        """
+        try:
+            agent_norm = agent_df.reset_index(drop=True).sort_index(axis=1)
+            expected_norm = expected_df.reset_index(drop=True).sort_index(axis=1)
+            
+            # Fast exact comparison
+            if set(agent_norm.columns) == set(expected_norm.columns):
+                
+                # Check directly first
+                if agent_norm.equals(expected_norm):
+                    return 1.0
+                
+                # Re-align column order if mismatched
+                agent_aligned = agent_norm[expected_norm.columns]
+                if agent_aligned.equals(expected_norm):
+                    return 1.0
+                
+                # Check unordered row set equivalence
+                agent_tuples = set(tuple(row) for row in agent_aligned.to_numpy())
+                expected_tuples = set(tuple(row) for row in expected_norm.to_numpy())
+                
+                if len(agent_tuples) == len(expected_tuples) and agent_tuples == expected_tuples:
+                    return 0.8  # Unordered match yields partial points
+            
+            return 0.0
+        except Exception:
+            return 0.0

core/tasks.py

@@ -0,0 +1,45 @@
+from dataclasses import dataclass
+from typing import List
+
+@dataclass
+class SqlTask:
+    difficulty: int
+    instruction: str
+    expected_query: str
+
+class EnvironmentTasks:
+    """Registry for all SQL challenges assigned to the OpenEnv agent."""
+    def __init__(self):
+        self.tasks: List[SqlTask] = [
+            SqlTask(
+                difficulty=1,
+                instruction="EASY: Find all names and salaries of employees, ordering by their salaries descending.",
+                expected_query="SELECT NAME, SALARY FROM EMPLOYEES ORDER BY SALARY DESC"
+            ),
+            SqlTask(
+                difficulty=2,
+                instruction="MEDIUM: Return the department names alongside the average salary per department. Columns must be 'DEPARTMENT_NAME' and 'AVG_SALARY'.",
+                expected_query="""
+                    SELECT D.NAME AS DEPARTMENT_NAME, AVG(E.SALARY) AS AVG_SALARY
+                    FROM EMPLOYEES E
+                    JOIN DEPARTMENTS D ON E.DEPARTMENT_ID = D.ID
+                    GROUP BY D.NAME
+                """
+            ),
+            SqlTask(
+                difficulty=3,
+                instruction="HARD: Find the name and title of all Engineering (DEPARTMENT_ID = 2) employees who earn a 'Senior' level salary (based on ROLES MIN_SALARY limit).",
+                expected_query="""
+                    SELECT E.NAME, R.TITLE 
+                    FROM EMPLOYEES E
+                    JOIN ROLES R ON E.SALARY >= R.MIN_SALARY
+                    WHERE E.DEPARTMENT_ID = 2 AND R.TITLE = 'Senior'
+                """
+            )
+        ]
+        
+    def get_total_tasks(self) -> int:
+        return len(self.tasks)
+        
+    def get_task(self, index: int) -> SqlTask:
+        return self.tasks[index]

database/__init__.py

@@ -0,0 +1,4 @@
+# Database package exports
+from database.sqlite_manager import SqliteManager
+
+__all__ = ["SqliteManager"]

database/sqlite_manager.py

@@ -0,0 +1,70 @@
+import sqlite3
+import pandas as pd
+from typing import List, Tuple
+
+class SqliteManager:
+    """Manages the lifecycle and execution of an internal SQLite database."""
+    def __init__(self, db_path: str = ":memory:"):
+        self.db_path = db_path
+        self.conn = None
+
+    def __enter__(self):
+        self.connect()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
+    def connect(self):
+        """Establishes the connection and sets up initial dummy schema + data."""
+        self.conn = sqlite3.Connection(self.db_path)
+        self._initialize_schema()
+
+    def close(self):
+        if self.conn:
+            self.conn.close()
+
+    def _initialize_schema(self):
+        cursor = self.conn.cursor()
+        
+        cursor.execute('''CREATE TABLE DEPARTMENTS (ID INTEGER PRIMARY KEY, NAME TEXT)''')
+        cursor.execute('''CREATE TABLE EMPLOYEES (ID INTEGER PRIMARY KEY, NAME TEXT, DEPARTMENT_ID INTEGER, SALARY REAL)''')
+        cursor.execute('''CREATE TABLE ROLES (ID INTEGER PRIMARY KEY, TITLE TEXT, MIN_SALARY REAL)''')
+        
+        cursor.executemany("INSERT INTO DEPARTMENTS (ID, NAME) VALUES (?, ?)", [
+            (1, "HR"), (2, "Engineering"), (3, "Marketing"), (4, "Sales")
+        ])
+        
+        cursor.executemany("INSERT INTO EMPLOYEES (ID, NAME, DEPARTMENT_ID, SALARY) VALUES (?, ?, ?, ?)", [
+            (1, "Alice", 2, 120000), (2, "Bob", 2, 110000), (3, "Charlie", 1, 80000),
+            (4, "David", 3, 90000), (5, "Eve", 4, 95000), (6, "Frank", 2, 105000),
+            (7, "Grace", 4, 100000), (8, "Hank", 2, 85000)
+        ])
+        
+        cursor.executemany("INSERT INTO ROLES (ID, TITLE, MIN_SALARY) VALUES (?, ?, ?)", [
+            (1, "Associate", 60000), (2, "Senior", 100000), (3, "Staff", 130000)
+        ])
+        
+        self.conn.commit()
+
+    def get_schema_summary(self) -> str:
+        return """
+Database Schema:
+1. DEPARTMENTS (ID INTEGER PRIMARY KEY, NAME TEXT)
+2. EMPLOYEES (ID INTEGER PRIMARY KEY, NAME TEXT, DEPARTMENT_ID INTEGER REFERENCES DEPARTMENTS(ID), SALARY REAL)
+3. ROLES (ID INTEGER PRIMARY KEY, TITLE TEXT, MIN_SALARY REAL)
+"""
+
+    def is_safe_query(self, query: str) -> bool:
+        """Validates that a query isn't destructive."""
+        forbidden = ["DROP", "DELETE", "UPDATE", "INSERT", "ALTER"]
+        upper_query = query.upper()
+        return not any(f in upper_query for f in forbidden)
+
+    def execute_dataframe(self, query: str) -> pd.DataFrame:
+        """Executes a SQL query and returns its result as a Pandas DataFrame."""
+        if not self.conn:
+            raise RuntimeError("Database connection not initialized.")
+        if not self.is_safe_query(query):
+            raise ValueError("Destructive operations are not allowed.")
+        return pd.read_sql_query(query, self.conn)

inference.py

@@ -0,0 +1,177 @@
+import os
+import textwrap
+from openai import OpenAI
+from typing import List
+from dotenv import load_dotenv
+
+load_dotenv()
+
+from client import SqlEnvClient
+from models import SqlAction
+
+# OpenEnv execution environment config
+ENV_CONTAINER_START = os.getenv("ENV_CONTAINER_START", "false").lower() == "true"
+ENV_IMAGE_NAME = "sql-agent-env:latest"
+
+# LLM inference config as per instructions
+API_BASE_URL = os.getenv("API_BASE_URL", "https://api.openai.com/v1")
+API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY") or os.getenv("OPENAI_API_KEY")
+GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
+MODEL_NAME = os.getenv("MODEL_NAME", "gpt-4o-mini")
+MAX_STEPS = 10
+
+SYSTEM_PROMPT = textwrap.dedent(
+    """
+    You are an expert Data Analyst and SQL Agent.
+    You will be provided with a SQL Schema and an instruction.
+    Reply with exactly one JSON block containing the SQL query to execute.
+    Do NOT provide markdown formatting or explanations. DO NOT wrap in ```sql.
+    The response should be the pure string of the query itself.
+    """
+).strip()
+
+def build_user_prompt(step: int, observation, history: List[str]) -> str:
+    schema = observation.schema_info
+    instruction = observation.current_task_instruction
+    exec_result = observation.execution_result or "None"
+    error = observation.execution_error or "None"
+    
+    prompt = textwrap.dedent(
+        f"""
+        Step: {step}
+        Database Schema:
+        {schema}
+        
+        Task Instruction:
+        {instruction}
+        
+        Previous Execution Result: {exec_result}
+        Previous Error: {error}
+        
+        Write the precise SQL query string to accomplish the task instruction. Return nothing else.
+        """
+    ).strip()
+    return prompt
+
+def parse_model_action(response_text: str) -> str:
+    query = response_text.strip()
+    if query.startswith("```sql"):
+        query = query[6:]
+    if query.startswith("```"):
+        query = query[3:]
+    if query.endswith("```"):
+        query = query[:-3]
+    return query.strip()
+
+def main():
+    if not API_KEY and not GEMINI_API_KEY:
+        print("WARNING: API credentials might be missing. Using dummy environment flow for validation.")
+    
+    # Client setup
+    client_type = "openai"
+    if GEMINI_API_KEY:
+        try:
+            from google import genai
+            client = genai.Client(api_key=GEMINI_API_KEY)
+            client_type = "gemini"
+            print("--- Using Google Gemini Client ---")
+        except ImportError:
+            print("WARNING: google-genai package not found. Falling back to OpenAI client.")
+            client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY or "dummy-key")
+    else:
+        client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY or "dummy-key")
+
+    print("--- SQL Data Analyst OpenEnv Baseline ---")
+    
+    if ENV_CONTAINER_START:
+        print(f"Starting docker image: {ENV_IMAGE_NAME}")
+        env = SqlEnvClient.from_docker_image(ENV_IMAGE_NAME).sync()
+    else:
+        # Fallback to local memory-bound initialization
+        from server.environment import SqlEnvironment
+        base_env = SqlEnvironment()
+        
+        # Direct wrapper
+        class DirectClient:
+            def __init__(self, target):
+                self.target = target
+            def reset(self):
+                obs = self.target.reset()
+                return type('StepResult', (), {'observation': obs, 'reward': 0.0, 'done': False})()
+            def step(self, action):
+                obs, reward, done = self.target.step(action)
+                return type('StepResult', (), {'observation': obs, 'reward': reward, 'done': done})()
+            def close(self):
+                pass
+        env = DirectClient(base_env)
+
+    history = []
+    
+    try:
+        result = env.reset()
+        observation = result.observation
+        print(f"\n[Episode Start] Initial Task: {observation.current_task_instruction}")
+        
+        for step in range(1, MAX_STEPS + 1):
+            if result.done:
+                print("Environment signalled done. Stopping early.")
+                break
+                
+            user_prompt = build_user_prompt(step, observation, history)
+            
+            messages = [
+                {"role": "system", "content": SYSTEM_PROMPT},
+                {"role": "user", "content": user_prompt},
+            ]
+            
+            try:
+                if client_type == "gemini":
+                    # For Gemini, we collapse System & User prompts into a single payload
+                    gemini_prompt = f"{SYSTEM_PROMPT}\n\n{user_prompt}"
+                    # Use provided model Name, fallback defaults to gemini-2.5-flash
+                    gemini_model = MODEL_NAME if MODEL_NAME != "gpt-4o-mini" else "gemini-2.5-flash"
+                    
+                    response = client.models.generate_content(
+                        model=gemini_model,
+                        contents=gemini_prompt,
+                        config={"temperature": 0.1}
+                    )
+                    response_text = response.text or "SELECT * FROM DEPARTMENTS"
+                elif API_KEY or os.getenv("OPENAI_API_KEY"):
+                    completion = client.chat.completions.create(
+                        model=MODEL_NAME,
+                        messages=messages,
+                        temperature=0.1,
+                        stream=False
+                    )
+                    response_text = completion.choices[0].message.content or "SELECT * FROM DEPARTMENTS"
+                else:
+                    response_text = "SELECT * FROM DEPARTMENTS"  # Dummy mode
+
+            except Exception as e:
+                print(f"Model Request failed: {e}. Using fallback action.")
+                response_text = "SELECT 1"
+
+            action_str = parse_model_action(response_text)
+            print(f"\n[Step {step}] Model suggested: {action_str}")
+            
+            result = env.step(SqlAction(query=action_str))
+            observation = result.observation
+            reward = result.reward
+            done = result.done
+
+            print(f"  Reward: {reward:+.2f} | Done: {done}")
+            if observation.execution_error:
+                print(f"  Error: {observation.execution_error[:200]}")
+                
+            history.append(f"Q: {action_str} -> R: {reward:+.2f}")
+            
+            if done:
+                print("\nAll tasks complete!")
+                break
+                
+    finally:
+        env.close()
+
+if __name__ == "__main__":
+    main()

models.py

@@ -0,0 +1,21 @@
+from pydantic import Field
+from typing import Optional, List
+from openenv.core.env_server import Action, Observation, State
+
+class SqlAction(Action):
+    """Action representing a SQL query execution."""
+    query: str
+
+class SqlObservation(Observation):
+    """Observation returned after executing a SQL query (or at reset)."""
+    current_task_instruction: str
+    schema_info: str
+    execution_result: Optional[str] = None
+    execution_error: Optional[str] = None
+
+class SqlState(State):
+    """State tracking for the SQL agent environment."""
+    current_task_index: int = 0
+    total_tasks: int = 3
+    accumulated_reward: float = 0.0
+    task_scores: List[float] = Field(default_factory=list)

openenv.yaml

@@ -0,0 +1,6 @@
+spec_version: 1
+name: sql-agent-env
+type: space
+runtime: fastapi
+app: server.app:app
+port: 8000

readme.md

@@ -0,0 +1,71 @@
+---
+title: SQL AI Data Analyst
+emoji: 📊
+colorFrom: blue
+colorTo: green
+sdk: docker
+app_port: 8000
+tags:
+  - openenv
+---
+# OpenEnv SQL Data Analyst
+
+## The Goal
+The objective is straightforward: the AI must interpret an internal database schema and translate human business instructions into valid SQL queries natively executed against the pipeline. 
+
+Tasks increase in difficulty, pushing models to understand complex relationships such as:
+- Simple Table Retrievals (`SELECT`, `WHERE`)
+- Aggregating related tables (`JOIN`, `GROUP BY`)
+- Complex constraint execution bounds (Logic matching)
+
+## Fast Execution
+1. **Install:** `pip install -r requirements.txt`
+2. **Setup your preferred engine:**
+   ```powershell
+   $env:GEMINI_API_KEY="your-api-key" 
+   # Or use OPENAI_API_KEY 
+   ```
+3. **Run the Simulation Base:**
+   ```powershell
+   python inference.py
+   ```
+
+## Architecture Workflow
+
+This environment follows strict [OpenEnv Gym interface](https://github.com/meta-pytorch/OpenEnv) standards, utilizing `step()`, `reset()`, and `state()` endpoints wrapped locally or exposed natively via Hugging Face.
+
+```mermaid
+flowchart TD
+    Start([Episode Start]) --> EnvInit[Environment loads Task 1]
+    
+    EnvInit --> Obs[Generate Observation]
+    Obs -.-> |1. Send Schema and Instruction| Agent((LLM Agent))
+    
+    Agent -.-> |2. Predicts SQL Query| Step[Environment step function]
+    
+    Step --> Exec{Execute SQL on SQLite Engine}
+    
+    Exec -->|SQL Syntax Error| Pen[Return Negative Reward]
+    Exec -->|Valid SQL| Grader[Evaluate Pandas DataFrame]
+    
+    Grader --> Match{Compare Expected DataFrame}
+    Match -->|Perfect Match| Rew1[Reward 1.0 Advance Level]
+    Match -->|Unordered Match| RewPartial[Reward 0.8]
+    Match -->|Incorrect Data| Rew0[Reward 0.0]
+    
+    Pen --> StateUpdate
+    Rew1 --> StateUpdate
+    RewPartial --> StateUpdate
+    Rew0 --> StateUpdate
+    
+    StateUpdate[Log History and Update Episode state] --> CheckDone{All Tasks Completed Yes or No}
+    
+    CheckDone -->|No| Obs
+    CheckDone -->|Yes| Finish([Episode End])
+```
+
+## Internal Engine Design
+To read more about exactly how the logic is handled under the hood, access the OOP components inside the root directory:
+- `server/environment.py`: The orchestrator handling the OpenEnv lifecycle.
+- `database/sqlite_manager.py`: Creates ephemeral SQLite sandboxes per session.
+- `core/grader.py`: Evaluates logic correctness independent of string semantics by leveraging Pandas DataFrame alignments.

requirements.txt

@@ -0,0 +1,7 @@
+fastapi>=0.111.0
+pydantic>=2.7.1
+uvicorn>=0.30.1
+openenv-core
+pandas
+google-genai
+python-dotenv

server/app.py

@@ -0,0 +1,7 @@
+from openenv.core.env_server import create_fastapi_app
+
+from models import SqlAction, SqlObservation
+from server.environment import SqlEnvironment
+
+env = SqlEnvironment()
+app = create_fastapi_app(env, SqlAction, SqlObservation)

server/environment.py

@@ -0,0 +1,95 @@
+import uuid
+import pandas as pd
+from typing import Tuple
+
+from openenv.core.env_server import Environment
+from models import SqlAction, SqlObservation, SqlState
+from database import SqliteManager
+from core import EnvironmentTasks, DataFrameGrader
+
+class SqlEnvironment(Environment):
+    """Orchestrates the SQL Agent execution matching the OpenEnv Gym API."""
+    def __init__(self):
+        super().__init__()
+        self._state = SqlState()
+        self.task_registry = EnvironmentTasks()
+        
+        # Instantiate db connection immediately for the server loop
+        self.db = SqliteManager()
+        self.db.connect()
+
+    def reset(self) -> SqlObservation:
+        self._state = SqlState(
+            episode_id=str(uuid.uuid4()),
+            current_task_index=0,
+            total_tasks=self.task_registry.get_total_tasks(),
+            accumulated_reward=0.0,
+            task_scores=[]
+        )
+        first_task = self.task_registry.get_task(0)
+        return SqlObservation(
+            current_task_instruction=first_task.instruction,
+            schema_info=self.db.get_schema_summary()
+        )
+
+    def state(self) -> SqlState:
+        return self._state
+
+    def step(self, action: SqlAction) -> Tuple[SqlObservation, float, bool]:
+        self._state.step_count += 1
+        query = action.query
+        current_idx = self._state.current_task_index
+        task = self.task_registry.get_task(current_idx)
+        
+        execution_result_str = None
+        error_str = None
+        reward = 0.0
+        done = False
+        
+        try:
+            # 1. Gather expected vs agent dataframes
+            expected_df = self.db.execute_dataframe(task.expected_query)
+            agent_df = self.db.execute_dataframe(query)
+            
+            # Serialize for observation
+            execution_result_str = agent_df.to_json(orient="records")
+            
+            # 2. Grade execution
+            score = DataFrameGrader.grade(agent_df, expected_df)
+            self._state.task_scores.append(score)
+            
+            # 3. Reward scaling
+            reward += 0.1 # Valid execution baseline
+            reward += score
+            
+            # 4. Check level progression
+            if score == 1.0:
+                 self._state.current_task_index += 1
+                 if self._state.current_task_index >= self.task_registry.get_total_tasks():
+                     done = True
+                     
+        except Exception as e:
+            error_str = str(e)
+            reward -= 0.5  # Heavy penalty for destructive/invalid SQL
+        
+        # Global accumulation
+        self._state.accumulated_reward += reward
+        
+        # Prepare observation response
+        if done:
+            total_score = sum(self._state.task_scores)
+            next_instruction = f"All tasks completed! Final SQL Capability Score: {total_score}"
+            schema = ""
+        else:
+            next_task = self.task_registry.get_task(self._state.current_task_index)
+            next_instruction = next_task.instruction
+            schema = self.db.get_schema_summary()
+            
+        obs = SqlObservation(
+            current_task_instruction=next_instruction,
+            schema_info=schema,
+            execution_result=execution_result_str,
+            execution_error=error_str
+        )
+        
+        return obs, reward, done