Upload 7 files

Dockerfile

@@ -0,0 +1,26 @@
+# Use a lightweight python image
+FROM python:3.11-slim
+
+# Install system dependencies (SQLite is built-in, but we need curl for healthchecks)
+RUN apt-get update && apt-get install -y curl && rm -rf /var/lib/apt/lists/*
+
+# Install uv (the fast python package manager)
+COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/
+
+# Set the working directory
+WORKDIR /app
+
+# Copy the dependency files first to cache the layer
+COPY pyproject.toml uv.lock ./
+
+# Install dependencies using uv
+RUN uv sync --frozen
+
+# Copy the rest of the application code
+COPY . .
+
+# Expose the port the FastAPI server runs on
+EXPOSE 8000
+
+# Start the OpenEnv server
+CMD ["uv", "run", "server/app.py"]

README.md

@@ -1,12 +1,18 @@
----
-title: Sqlite Rescue Env
-emoji: 📉
-colorFrom: gray
-colorTo: pink
-sdk: docker
-pinned: false
-license: mit
-short_description: AI agents rescue, clean, and normalize a messy SQLite dbs
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# SQLite Rescue Environment
+
+## Description
+A data engineering environment where an AI agent must clean, normalize, and manipulate a messy SQLite database using raw SQL queries. This simulates the real-world task of data cleaning and database refactoring.
+
+## Action & Observation Spaces
+* **Action Space:** The agent submits an SQL `query` (string) to execute, and a `submit` (boolean) flag when they are ready for their final database state to be graded.
+* **Observation Space:** The environment returns the current `schema_info` (string), `rows_affected` (int), any SQL execution `error` (string), and a `query_result` (list of dicts) if the action was a SELECT query.
+
+## Tasks
+1. **easy_data_cleaning:** Clean inconsistent dates and trailing whitespaces in a single table.
+2. **medium_schema_normalization:** Split a denormalized monolithic table into two related tables with a foreign key.
+3. **hard_complex_reconciliation:** Write a complex query/view to generate a financial reconciliation report.
+
+## Setup Instructions
+1. Install dependencies: `pip install -r requirements.txt` (or via pyproject.toml)
+2. Generate the starting databases: `python generate_templates.py`
+3. Run the baseline: `python inference.py`

env.py

@@ -0,0 +1,98 @@
+import sqlite3
+import shutil
+import os
+from typing import Tuple, Dict, Any
+
+from models import RescueAction, RescueObservation, RescueState
+from graders import grade_easy_task, grade_medium_task, grade_hard_task
+
+class DatabaseRescueEnv:
+    def __init__(self):
+        self.task_name = None
+        self.steps_taken = 0
+        self.working_db = "working.db"
+        self.template_dir = "templates"
+
+    def _get_schema(self, cursor: sqlite3.Cursor) -> str:
+        """Retrieves the current database schema as a string."""
+        cursor.execute("SELECT sql FROM sqlite_master WHERE type='table';")
+        tables = cursor.fetchall()
+        return "\n".join([t[0] for t in tables if t[0] is not None])
+
+    def reset(self, task_name: str) -> RescueObservation:
+        """Resets the environment by copying the task's template DB."""
+        self.task_name = task_name
+        self.steps_taken = 0
+        
+        # Copy the messy template DB to our working path
+        template_path = os.path.join(self.template_dir, f"{task_name}.db")
+        if not os.path.exists(template_path):
+            raise ValueError(f"Template DB for task '{task_name}' not found.")
+            
+        shutil.copyfile(template_path, self.working_db)
+        
+        # Return initial observation
+        with sqlite3.connect(self.working_db) as conn:
+            schema = self._get_schema(conn.cursor())
+            
+        return RescueObservation(
+            schema_info=schema,
+            query_result=None,
+            rows_affected=0,
+            error=None
+        )
+
+    def step(self, action: RescueAction) -> Tuple[RescueObservation, float, bool, Dict[str, Any]]:
+        self.steps_taken += 1
+        reward = 0.0
+        done = False
+        info = {}
+
+        # If the agent submits, trigger the grader and end the episode
+        if action.submit:
+            reward = self._grade_task()
+            done = True
+            with sqlite3.connect(self.working_db) as conn:
+                schema = self._get_schema(conn.cursor())
+            return RescueObservation(schema_info=schema), reward, done, info
+
+        # Otherwise, execute the query
+        obs = RescueObservation(schema_info="", rows_affected=0)
+        
+        try:
+            with sqlite3.connect(self.working_db) as conn:
+                conn.row_factory = sqlite3.Row 
+                cursor = conn.cursor()
+                
+                cursor.execute(action.query)
+                conn.commit()
+                
+                obs.schema_info = self._get_schema(cursor)
+                obs.rows_affected = cursor.rowcount
+                
+                # If it was a SELECT query, fetch results
+                if action.query.strip().upper().startswith("SELECT"):
+                    rows = cursor.fetchmany(50) # Limit to avoid context bloat
+                    obs.query_result = [dict(row) for row in rows]
+                    
+        except sqlite3.Error as e:
+            obs.error = str(e)
+            
+        return obs, reward, done, info
+
+    def state(self) -> RescueState:
+        return RescueState(
+            task_name=self.task_name,
+            steps_taken=self.steps_taken,
+            db_path=self.working_db
+        )
+        
+    def _grade_task(self) -> float:
+        """Routes the current database to the correct scoring logic."""
+        if self.task_name == "easy_data_cleaning":
+            return grade_easy_task(self.working_db)
+        elif self.task_name == "medium_schema_normalization":
+            return grade_medium_task(self.working_db)
+        elif self.task_name == "hard_complex_reconciliation":
+            return grade_hard_task(self.working_db)
+        return 0.0

inference.py

@@ -0,0 +1,74 @@
+import os
+from openai import OpenAI
+from env import DatabaseRescueEnv
+from models import RescueAction
+
+# --- CONFIGURATION (Mapped to Hackathon Requirements) ---
+API_KEY = os.getenv("HF_TOKEN") or os.getenv("OPENAI_API_KEY")
+API_BASE_URL = os.getenv("API_BASE_URL") or "https://router.huggingface.co/v1"
+MODEL_NAME = os.getenv("MODEL_NAME") or "Qwen/Qwen2.5-72B-Instruct"
+TASK_NAME = "easy_data_cleaning"
+MAX_STEPS = 5
+
+def run_baseline():
+    # Initialize the client and the environment
+    client = OpenAI(api_key=API_KEY, base_url=API_BASE_URL)
+    env = DatabaseRescueEnv()
+    
+    # 1. Print the mandatory [START] log
+    print(f"[START] task={TASK_NAME} env=sqlite-rescue-env model={MODEL_NAME}")
+    
+    # Reset the environment
+    obs = env.reset(TASK_NAME)
+    rewards = []
+    success = False
+    error_msg = "null"
+    
+    # The hardcoded solution to the Easy task for the baseline agent
+    solution_queries = [
+        "UPDATE customers SET name = TRIM(name);",
+        "UPDATE customers SET signup_date = substr(signup_date, 7, 4) || '-' || substr(signup_date, 1, 2) || '-' || substr(signup_date, 4, 2) WHERE signup_date LIKE '%/%';",
+        "UPDATE customers SET signup_date = substr(signup_date, 7, 4) || '-' || substr(signup_date, 1, 2) || '-' || substr(signup_date, 4, 2) WHERE signup_date LIKE '%-%' AND length(signup_date) = 10 AND substr(signup_date, 3, 1) = '-';",
+        "SELECT * FROM customers;"
+    ]
+
+    steps_taken = 0
+    for i in range(MAX_STEPS):
+        steps_taken += 1
+        
+        # Decide the action
+        if i < len(solution_queries):
+            query = solution_queries[i]
+            action = RescueAction(query=query, submit=False)
+            action_str = f"execute_sql('{query}')"
+        else:
+            action = RescueAction(query="", submit=True)
+            action_str = "submit(True)"
+
+        # Step the environment
+        obs, reward, done, info = env.step(action)
+        rewards.append(reward)
+        
+        if obs.error:
+            error_msg = f"'{obs.error}'"
+        else:
+            error_msg = "null"
+
+        # 2. Print the mandatory [STEP] log
+        print(f"[STEP] step={steps_taken} action={action_str} reward={reward:.2f} done={str(done).lower()} error={error_msg}")
+
+        if done:
+            success = (reward == 1.0)
+            break
+            
+    # 3. Print the mandatory [END] log
+    rewards_str = ",".join([f"{r:.2f}" for r in rewards])
+    final_score = rewards[-1] if rewards else 0.00
+    print(f"[END] success={str(success).lower()} steps={steps_taken} score={final_score:.2f} rewards={rewards_str}")
+
+if __name__ == "__main__":
+    # Ensure the user has an API key set
+    if not API_KEY:
+        print("Error: Please set HF_TOKEN or OPENAI_API_KEY environment variable.")
+    else:
+        run_baseline()

models.py

@@ -0,0 +1,17 @@
+from pydantic import BaseModel, Field
+from typing import Optional, List, Dict, Any
+
+class RescueAction(BaseModel):
+    query: str = Field(description="The SQL query to execute against the SQLite database.")
+    submit: bool = Field(default=False, description="Set to True ONLY when you have completed the task and are ready for grading.")
+
+class RescueObservation(BaseModel):
+    schema_info: str = Field(description="The current schema of the database.")
+    query_result: Optional[List[Dict[str, Any]]] = Field(default=None, description="Results from a SELECT query, limited to 50 rows.")
+    rows_affected: int = Field(default=0, description="Number of rows modified by INSERT/UPDATE/DELETE.")
+    error: Optional[str] = Field(default=None, description="SQL execution error message, if any.")
+
+class RescueState(BaseModel):
+    task_name: str
+    steps_taken: int
+    db_path: str

pyproject.toml

@@ -0,0 +1,20 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "sqlite-rescue-env"
+version = "1.0.0"
+description = "A data engineering environment where an agent must rescue a messy SQLite database."
+readme = "README.md"
+requires-python = ">=3.9"
+dependencies = [
+    "openenv-core>=0.1.0",
+    "pydantic>=2.0.0",
+    "openai>=1.0.0",
+    "fastapi>=0.100.0",
+    "uvicorn>=0.23.0"
+]
+
+[project.scripts]
+server = "server.app:main"