server/app.py

#!/usr/bin/env python3
"""
OpenEnv Server Entry Point for Data Cleaning Environment.
This module provides the server entry point for multi-mode deployment.
"""

import sys
import os

# Add parent directory to path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from typing import Optional, Dict, Any
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import uvicorn

from env.environment import DataCleaningEnvironment

# Create FastAPI app
app = FastAPI(
    title="Data Cleaning OpenEnv API",
    description="OpenEnv environment for data cleaning tasks",
    version="1.0.0"
)

# Add CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Global environment instance
env_instance: Optional[DataCleaningEnvironment] = None
current_obs = None


# ============== API Models ==============

class ResetRequest(BaseModel):
    task: str = "easy"
    seed: Optional[int] = 42

class ActionRequest(BaseModel):
    action_type: str
    column: Optional[str] = None
    fill_strategy: Optional[str] = None
    fill_value: Optional[Any] = None
    normalize_method: Optional[str] = None
    text_case: Optional[str] = None
    outlier_threshold: Optional[float] = None


# ============== API Endpoints ==============

@app.get("/")
async def root():
    """Health check endpoint."""
    return {"status": "ok", "environment": "data-cleaning-openenv", "version": "1.0.0"}


@app.post("/reset")
async def api_reset(request: ResetRequest = None):
    """Reset the environment and return initial observation."""
    global env_instance, current_obs
    
    if request is None:
        request = ResetRequest()
    
    try:
        task = request.task if request.task in ["easy", "medium", "hard"] else "easy"
        seed = request.seed if request.seed is not None else 42
        
        env_instance = DataCleaningEnvironment(task=task, seed=seed)
        current_obs = env_instance.reset()
        
        return {
            "observation": current_obs.model_dump(),
            "info": {
                "task": task,
                "seed": seed
            }
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/reset")
async def api_reset_get(task: str = "easy", seed: int = 42):
    """Reset the environment (GET version for compatibility)."""
    return await api_reset(ResetRequest(task=task, seed=seed))


@app.post("/step")
async def api_step(action: ActionRequest):
    """Execute an action and return the result."""
    global env_instance, current_obs
    
    if env_instance is None:
        await api_reset(ResetRequest())
    
    try:
        action_dict = {"action_type": action.action_type}
        
        if action.column:
            action_dict["column"] = action.column
        if action.fill_strategy:
            action_dict["fill_strategy"] = action.fill_strategy
        if action.fill_value is not None:
            action_dict["fill_value"] = action.fill_value
        if action.normalize_method:
            action_dict["normalize_method"] = action.normalize_method
        if action.text_case:
            action_dict["text_case"] = action.text_case
        if action.outlier_threshold is not None:
            action_dict["outlier_threshold"] = action.outlier_threshold
        
        result = env_instance.step(action_dict)
        current_obs = result.observation
        
        return {
            "observation": result.observation.model_dump(),
            "reward": result.reward.model_dump() if result.reward else {"total": 0},
            "done": result.done,
            "info": result.info
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/state")
async def api_state():
    """Get current environment state."""
    global env_instance
    
    if env_instance is None:
        raise HTTPException(status_code=400, detail="Environment not initialized. Call /reset first.")
    
    try:
        state = env_instance.state()
        return {"state": state.model_dump()}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/tasks")
async def api_tasks():
    """List available tasks."""
    return {
        "tasks": [
            {"id": "easy", "description": "Simple dataset with missing values and duplicates", "max_steps": 15},
            {"id": "medium", "description": "Medium complexity with formatting issues", "max_steps": 25},
            {"id": "hard", "description": "Complex dataset with outliers and mixed issues", "max_steps": 40}
        ]
    }


@app.get("/actions")
async def api_actions():
    """List available actions."""
    return {
        "actions": [
            {"type": "fill_missing", "params": ["column", "fill_strategy", "fill_value"]},
            {"type": "remove_duplicates", "params": []},
            {"type": "normalize_column", "params": ["column", "normalize_method"]},
            {"type": "detect_outliers", "params": ["column", "outlier_threshold"]},
            {"type": "standardize_text", "params": ["column", "text_case"]},
            {"type": "drop_column", "params": ["column"]},
            {"type": "finish", "params": []}
        ]
    }


def main():
    """Main entry point for the server."""
    uvicorn.run(app, host="0.0.0.0", port=7860)


if __name__ == "__main__":
    main()