import uvicorn
import sys
import os
import sqlite3
from fastapi import FastAPI
from fastapi.responses import HTMLResponse
# Force the root directory into the python path
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
if parent_dir not in sys.path:
sys.path.insert(0, parent_dir)
import env
app = FastAPI(title="SQLite Rescue Environment API")
env_instance = env.DatabaseRescueEnv()
# @app.get("/")
# def health_check():
# return {"status": "ok", "environment": "sqlite-rescue-env"}
@app.get("/", response_class=HTMLResponse)
def serve_dashboard():
"""Generates a live HTML view of the current database."""
# IMPORTANT: Change "workspace.db" if your environment uses a different filename!
db_path = "workspace.db"
html_content = """
🗄️ SQLite Rescue - Live View
🗄️ SQLite Rescue Environment - Live View
This dashboard shows the real-time state of the agent's workspace database. Refresh the page to see changes after an agent takes an action.
"""
if not os.path.exists(db_path):
html_content += f"Waiting for environment to initialize... ({db_path} not found).
"
else:
try:
with sqlite3.connect(db_path) as conn:
c = conn.cursor()
# Find all tables and views in the database
c.execute("SELECT name, type FROM sqlite_master WHERE type IN ('table', 'view') AND name NOT LIKE 'sqlite_%';")
db_objects = c.fetchall()
if not db_objects:
html_content += "Database is empty.
"
for obj_name, obj_type in db_objects:
html_content += f""
html_content += f"
{obj_name} ({obj_type})
"
# Fetch all rows for the table
c.execute(f"SELECT * FROM {obj_name} LIMIT 50")
rows = c.fetchall()
if not rows:
html_content += "
No data
"
for col in cols:
html_content += f"| {col} | "
html_content += "
"
for row in rows:
html_content += ""
for val in row:
html_content += f"| {str(val)} | "
html_content += "
"
html_content += "
"
except Exception as e:
html_content += f"Error reading database: {str(e)}
"
html_content += ""
return HTMLResponse(content=html_content, status_code=200)
@app.post("/reset")
def reset_env(task_name: str = "easy_data_cleaning"):
obs = env_instance.reset(task_name)
return {"status": "reset", "observation": obs.model_dump()} # Changed .dict() to .model_dump() for Pydantic v2
def main():
uvicorn.run(app, host="0.0.0.0", port=7860) # Port 7860 is the HF default
if __name__ == "__main__":
main()