server/app.py

"""
FastAPI application for the SRE OpenEnv environment.

Creates the HTTP server that exposes the canonical OpenEnv endpoints
(reset, step, state, health) via the create_fastapi_app helper.
"""

from openenv_core.env_server import create_fastapi_app

from models import SREAction, SREObservation, to_dict
from server.sre_environment import SREEnvironment

from fastapi import FastAPI, Body
from typing import Dict, Any

# Initialize environment
env_instance = SREEnvironment()
app = FastAPI(title="SRE OpenEnv Server")

@app.post("/reset")
async def reset(request: Dict[str, Any] = Body(default={})):
    """Reset with support for task_id and seed."""
    seed = request.get("seed")
    task_id = request.get("task_id")
    
    # env_instance.reset handles both if provided
    observation = env_instance.reset(seed=seed, task_id=task_id)
    
    # Serialize to OpenEnv StepResult format
    obs_dict = to_dict(observation)
    reward = obs_dict.pop("reward", 0.0)
    done = obs_dict.pop("done", False)
    
    return {
        "observation": obs_dict,
        "reward": reward,
        "done": done,
    }

@app.post("/step")
async def step(request: Dict[str, Any]):
    """Step with SREAction deserialization."""
    action_data = request.get("action", {})
    action = SREAction(**action_data)
    observation = env_instance.step(action)
    
    obs_dict = to_dict(observation)
    reward = obs_dict.pop("reward", 0.0)
    done = obs_dict.pop("done", False)
    
    return {
        "observation": obs_dict,
        "reward": reward,
        "done": done,
    }

@app.get("/state")
async def get_state():
    return to_dict(env_instance.state)

@app.get("/health")
async def health():
    return {"status": "healthy"}


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


if __name__ == "__main__":
    main()