minimal_openenv_app.py

from __future__ import annotations

import os
import time
from pathlib import Path
from typing import Literal
from uuid import uuid4

from fastapi import Request
from fastapi.responses import PlainTextResponse, RedirectResponse
from openenv.core.env_server import create_app
from openenv.core.env_server.interfaces import Environment
from openenv.core.env_server.types import Action, Observation, State
from pydantic import Field


LOG_PATH = Path("/tmp/minimal-space.log")
START_TS = time.strftime("%Y-%m-%dT%H:%M:%S%z")


def log(message: str) -> None:
    line = f"[{time.strftime('%Y-%m-%dT%H:%M:%S%z')}] {message}"
    print(line, flush=True)
    try:
        with LOG_PATH.open("a", encoding="utf-8") as f:
            f.write(line + "\n")
    except Exception:
        pass


class MinimalAction(Action):
    action_type: Literal["noop", "increment", "finish"] = "noop"
    amount: int = Field(default=1, ge=1, le=3)


class MinimalObservation(Observation):
    status: str
    counter: int
    summary: str
    reward: float = 0.0
    done: bool = False


class MinimalState(State):
    counter: int = 0


class MinimalEnvironment(Environment[MinimalAction, MinimalObservation, MinimalState]):
    SUPPORTS_CONCURRENT_SESSIONS = False

    def __init__(self):
        super().__init__()
        log("MinimalEnvironment.__init__ begin")
        self._done = False
        self._state = MinimalState(episode_id=str(uuid4()), step_count=0, counter=0)
        log("MinimalEnvironment.__init__ end")

    def reset(self, seed: int | None = None, episode_id: str | None = None, **kwargs) -> MinimalObservation:
        log(f"reset begin seed={seed} episode_id={episode_id} kwargs={kwargs}")
        self._done = False
        self._state = MinimalState(
            episode_id=episode_id or str(uuid4()),
            step_count=0,
            counter=0,
        )
        obs = self._observation(status="ready", reward=0.0, done=False)
        log(f"reset end episode_id={self._state.episode_id}")
        return obs

    def step(self, action: MinimalAction, timeout_s: float | None = None, **kwargs) -> MinimalObservation:
        log(f"step begin action={action.model_dump()} timeout_s={timeout_s} kwargs={kwargs}")
        if self._done:
            obs = self._observation(status="done", reward=0.0, done=True)
            log("step end already done")
            return obs

        self._state.step_count += 1
        reward = 0.0
        status = "ok"

        if action.action_type == "increment":
            self._state.counter += action.amount
            reward = float(action.amount)
        elif action.action_type == "finish":
            self._done = True
            status = "finished"

        if self._state.step_count >= 8:
            self._done = True
            status = "finished"

        obs = self._observation(status=status, reward=reward, done=self._done)
        log(f"step end counter={self._state.counter} step_count={self._state.step_count} done={self._done}")
        return obs

    @property
    def state(self) -> MinimalState:
        log("state property")
        return self._state

    def close(self) -> None:
        log("close")

    def _observation(self, *, status: str, reward: float, done: bool) -> MinimalObservation:
        return MinimalObservation(
            status=status,
            counter=self._state.counter,
            summary=(
                f"Minimal OpenEnv demo. Counter={self._state.counter}. "
                f"Step={self._state.step_count}. "
                f"Choose noop, increment, or finish."
            ),
            reward=reward,
            done=done,
        )


log("minimal_openenv_app module import begin")
app = create_app(MinimalEnvironment, MinimalAction, MinimalObservation, env_name="minimal_openenv")
log("openenv app created")


@app.on_event("startup")
async def startup() -> None:
    log("startup event begin")
    log(f"python={os.sys.version.split()[0]}")
    log(f"cwd={os.getcwd()}")
    log(f"enable_web={os.getenv('ENABLE_WEB_INTERFACE')}")
    log("startup event end")


@app.middleware("http")
async def request_logger(request: Request, call_next):
    log(f"request start method={request.method} path={request.url.path}")
    response = await call_next(request)
    log(f"request end method={request.method} path={request.url.path} status={response.status_code}")
    return response


@app.get("/", include_in_schema=False)
def root() -> RedirectResponse:
    log("root redirect")
    return RedirectResponse(url="/web")


@app.get("/logs", include_in_schema=False)
def logs() -> PlainTextResponse:
    log("logs handler")
    try:
        return PlainTextResponse(LOG_PATH.read_text(encoding="utf-8"))
    except FileNotFoundError:
        return PlainTextResponse("no log file yet\n")


log("minimal_openenv_app module import end")