org_sim/server/app.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
# BSD-style license

"""FastAPI application for the OrgSim Environment."""

import os
from fastapi import FastAPI
from openenv.core.env_server.http_server import create_app

from ..models import OrgAction, OrgObservation, OrgState
from .org_environment import OrgSimEnvironment, TASK_IDS

_env: OrgSimEnvironment | None = None


def get_env() -> OrgSimEnvironment:
    global _env
    if _env is None:
        max_steps = int(os.getenv("ORGSIM_MAX_STEPS", "50"))
        _env = OrgSimEnvironment(max_steps=max_steps)
    return _env


def create_orgsim_app() -> FastAPI:
    """Create the OrgSim FastAPI app with /grade and /tasks endpoints."""
    app = create_app(
        get_env,
        OrgAction,
        OrgObservation,
        env_name="org_sim",
    )

    @app.get("/tasks")
    def list_tasks():
        """List available task IDs for enumeration by automated validator."""
        return {"tasks": TASK_IDS}

    @app.get("/grade")
    def grade_episode():
        """Return current episode grade (call after done=True)."""
        env = get_env()
        score = env.grade_episode()
        return {"score": score, "task_id": env._current_task_id}

    return app


app = create_orgsim_app()


def main():
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)


if __name__ == "__main__":
    main()