server/app.py

"""FastAPI application for the BrainRL region selection environment.

Three concerns live in this module on purpose:

1. ``create_app(...)`` exposes the OpenEnv MCP/HTTP wire format for the BrainRL
   region-selection environment. This is the contract Colab/HF Jobs talk to.
2. The trained-policy service (``LazyTransformersPolicy``) lets a GPU Space
   load a pushed checkpoint and emit actions on demand.
3. Operational endpoints (``/health``, ``/brainrl/data_status``,
   ``/brainrl/manifest``) let any caller verify *which* data revision and
   *which* model revision the Space is actually serving. Without these you
   cannot tell from the outside whether a redeploy actually picked up the new
   checkpoint, which is the bug the plan's "end-to-end-validation" item is
   trying to prevent.
"""

import os
from typing import Any

try:
    from hf_data import parcel_manifest_summary, sync_data_from_env
except ImportError:  # pragma: no cover
    sync_data_from_env = None  # type: ignore[assignment]
    parcel_manifest_summary = None  # type: ignore[assignment]

_DATA_STATUS: dict[str, Any] | None = sync_data_from_env() if sync_data_from_env else None

try:
    from openenv.core.env_server.http_server import create_app
    from openenv.core.env_server.mcp_types import CallToolAction, CallToolObservation
    from .brain_environment import BrainRegionSelectionEnvironment
    from .policy_service import LazyTransformersPolicy
    from .prediction_ui import build_prediction_ui
    from prompts import parse_region_action
except ImportError:
    from openenv.core.env_server.http_server import create_app
    from openenv.core.env_server.mcp_types import CallToolAction, CallToolObservation
    from server.brain_environment import BrainRegionSelectionEnvironment
    from server.policy_service import LazyTransformersPolicy
    from server.prediction_ui import build_prediction_ui
    from prompts import parse_region_action


_POLICY = LazyTransformersPolicy()


def run_brain_rollout(
    *,
    seed: int = 42,
    subject_id: str | None = None,
    run_id: str | None = None,
    condition: str | None = None,
    stimulus_window: int | None = None,
) -> dict[str, Any]:
    """Run one trained-policy episode and return rollout metadata.

    Shared between the ``/rollout`` HTTP endpoint and the Gradio prediction UI
    so both surfaces always see identical step traces. Each step records the
    parcel chosen plus the slice of candidate metadata (label, network,
    hemisphere) the model could see at decision time, so the UI can render
    "what did the model just predict, and why".
    """

    env = BrainRegionSelectionEnvironment()
    env.reset(
        seed=int(seed),
        subject_id=subject_id,
        run_id=run_id,
        condition=condition,
        stimulus_window=stimulus_window,
    )

    steps: list[dict[str, Any]] = []
    done = False
    max_steps = int(env._subset.selection_budget) + 2  # belt-and-braces guard
    while not done and len(steps) < max_steps:
        state = env._build_selection_state()
        candidate_index = {
            str(item["region_id"]): item for item in state.get("candidate_regions", [])
        }
        action = _POLICY.choose_action(state)
        chosen = candidate_index.get(str(action.region_id), {})
        result = env._process_action(action.region_id)
        previous_r2 = float(result.get("previous_r2", 0.0))
        current_r2 = float(result["current_r2"])
        steps.append(
            {
                "timestep": env._timestep,
                "region_id": action.region_id,
                "region_label": chosen.get("label"),
                "atlas": chosen.get("atlas"),
                "hemisphere": chosen.get("hemisphere"),
                "network": chosen.get("network"),
                "sub_region": chosen.get("sub_region"),
                "redundancy_group": chosen.get("redundancy_group"),
                "base_r2": chosen.get("base_r2"),
                "semantic_prior": chosen.get("semantic_prior"),
                "reward": float(result["reward"]),
                "current_r2": current_r2,
                "delta_r2": current_r2 - previous_r2,
                "done": bool(result["done"]),
                "error": result.get("error"),
                "policy_error": action.error,
                "fallback_used": action.fallback_used,
                "raw_text": action.raw_text,
            }
        )
        done = bool(result["done"])

    final_state = env._build_selection_state()
    return {
        "subject_id": env._subject_id,
        "run_id": env._run_id,
        "condition": env._condition,
        "final_r2": float(env._current_r2),
        "selection_budget": int(env._subset.selection_budget),
        "selected_regions": list(env._selected_region_ids),
        "stimulus": final_state.get("stimulus"),
        "stimulus_window": final_state.get("stimulus_window"),
        "stimulus_n_windows": final_state.get("stimulus_n_windows"),
        "candidate_count": int(env._subset.n_regions),
        "atlas": env._subset.atlas,
        "steps": steps,
        "policy": _POLICY.status(),
    }


app = create_app(
    BrainRegionSelectionEnvironment,
    CallToolAction,
    CallToolObservation,
    env_name="brain_rl_region_selection",
    gradio_builder=build_prediction_ui(run_brain_rollout, _POLICY),
)


def _coerce_stimulus_window(value: Any) -> int | None:
    """Accept None / int / numeric string and reject negative values."""

    if value is None or value == "":
        return None
    try:
        coerced = int(value)
    except (TypeError, ValueError):
        return None
    return coerced if coerced >= 0 else None


def _empty_data_status() -> dict[str, Any]:
    """Build the data-status payload when no HF Dataset sync has run.

    Even in this mode (CPU/dev Space that ships configs in the image) we want
    the parcel manifest checksum reported, otherwise ``/brainrl/health`` will
    falsely flag the Space as broken.
    """

    config_dir = os.getenv("BRAINRL_CONFIG_DIR")
    manifest = {"present": False}
    if config_dir and parcel_manifest_summary is not None:
        from pathlib import Path

        manifest = parcel_manifest_summary(Path(config_dir) / "parcel_candidates.json")
    return {
        "root": None,
        "config_dir": config_dir,
        "has_parcel_manifest": bool(manifest.get("present")),
        "parcel_manifest": manifest,
        "annotation_csv_count": None,
        "data_repo": os.getenv("BRAINRL_DATA_REPO"),
        "data_revision": os.getenv("BRAINRL_DATA_REVISION") or None,
    }


@app.get("/brainrl/data_status")
def data_status() -> dict[str, Any]:
    """Return the HF Dataset/config status used by this Space."""

    return _DATA_STATUS or _empty_data_status()


@app.get("/brainrl/manifest")
def brainrl_manifest() -> dict[str, Any]:
    """Return parcel manifest checksum + budget so callers can pin revisions."""

    if _DATA_STATUS and _DATA_STATUS.get("parcel_manifest"):
        return _DATA_STATUS["parcel_manifest"]
    if parcel_manifest_summary is None:
        return {"present": False}
    config_dir = os.getenv("BRAINRL_CONFIG_DIR")
    if not config_dir:
        return {"present": False}
    from pathlib import Path

    return parcel_manifest_summary(Path(config_dir) / "parcel_candidates.json")


@app.get("/policy/status")
def policy_status() -> dict[str, Any]:
    """Return trained-policy loading status for GPU Spaces."""

    return _POLICY.status()


@app.get("/brainrl/health")
def brainrl_health() -> dict[str, Any]:
    """End-to-end readiness probe for the Space.

    Aggregates the three things that can independently silently break:
    the HF Dataset sync, the trained-policy load, and the environment's
    ability to actually build a state. The final ``ok`` is a conjunction
    so a green response is sufficient evidence to start a rollout.

    Note: we cannot use ``/health`` because OpenEnv's ``create_app`` already
    registers a liveness probe there that just returns ``{"status":"healthy"}``.
    """

    data = _DATA_STATUS or _empty_data_status()
    policy = _POLICY.status()

    env_ready = False
    env_error: str | None = None
    env_summary: dict[str, Any] = {}
    try:
        probe = BrainRegionSelectionEnvironment()
        probe.reset(seed=0)
        state = probe._build_selection_state()
        env_summary = {
            "selection_budget": int(state.get("selection_budget", 0)),
            "candidate_count": len(state.get("candidate_regions", [])),
            "subject_id": state.get("subject_id"),
            "condition": state.get("condition"),
        }
        env_ready = env_summary["candidate_count"] > 0
    except Exception as exc:  # pragma: no cover - report instead of 500
        env_error = f"{type(exc).__name__}: {exc}"

    data_ok = bool(data.get("config_dir")) and bool(
        data.get("parcel_manifest", {}).get("present")
    )
    policy_ok = (not policy["enabled"]) or (policy["loaded"] or policy["load_error"] is None)
    ok = bool(env_ready and data_ok and policy_ok)

    return {
        "ok": ok,
        "data": data,
        "policy": policy,
        "environment": {"ready": env_ready, "error": env_error, **env_summary},
        "space": {
            "config_dir": os.getenv("BRAINRL_CONFIG_DIR"),
            "stimulus_dir": os.getenv("BRAINRL_STIMULUS_DIR"),
            "data_repo": os.getenv("BRAINRL_DATA_REPO"),
            "data_revision": os.getenv("BRAINRL_DATA_REVISION") or None,
        },
    }


@app.post("/policy/action")
def policy_action(payload: dict[str, Any]) -> dict[str, Any]:
    """Choose one parcel action from a BrainRL selection state."""

    state = payload.get("state", payload)
    if not isinstance(state, dict):
        raise ValueError("Expected request body to contain a BrainRL state dict.")
    action = _POLICY.choose_action(state)
    return {
        "region_id": action.region_id,
        "raw_text": action.raw_text,
        "error": action.error,
        "fallback_used": action.fallback_used,
        "policy": _POLICY.status(),
    }


@app.post("/verifier/score_action")
def verifier_score_action(payload: dict[str, Any]) -> dict[str, Any]:
    """Score one model completion in a (possibly partial) BrainRL state.

    When the caller passes ``prior_selections`` we replay those parcels on a
    fresh env after reset, then score the model's action against the resulting
    intermediate state. This mirrors the local replay path in
    ``train_grpo.verify_completion`` so remote-env training (``--env-url``)
    produces the same rewards as in-process training.
    """

    env = BrainRegionSelectionEnvironment()
    env.reset(
        seed=int(payload.get("seed", 0)),
        subject_id=payload.get("subject_id"),
        run_id=payload.get("run_id"),
        condition=payload.get("condition"),
        stimulus_window=_coerce_stimulus_window(payload.get("stimulus_window")),
    )
    prior_selections = payload.get("prior_selections") or []
    if not isinstance(prior_selections, list):
        prior_selections = []
    for prior_id in prior_selections:
        prior_id_str = str(prior_id)
        if prior_id_str:
            env._process_action(prior_id_str)

    state = env._build_selection_state()
    valid_region_ids = {
        str(item["region_id"])
        for item in state["candidate_regions"]
        if not bool(item.get("selected"))
    }
    completion = str(payload.get("completion", ""))
    region_id, parse_error = parse_region_action(completion, valid_region_ids)
    if parse_error or region_id is None:
        result = env._process_action("__invalid__")
        result["parse_error"] = parse_error
        return result
    return env._process_action(region_id)


@app.post("/rollout")
def rollout(payload: dict[str, Any]) -> dict[str, Any]:
    """Run a full model-controlled rollout inside the Space.

    Thin wrapper around :func:`run_brain_rollout`; the helper is also driven by
    the Gradio prediction UI at ``/web`` so both surfaces report identical
    step traces.
    """

    return run_brain_rollout(
        seed=int(payload.get("seed", 42)),
        subject_id=payload.get("subject_id"),
        run_id=payload.get("run_id"),
        condition=payload.get("condition"),
        stimulus_window=_coerce_stimulus_window(payload.get("stimulus_window")),
    )


def main():
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)


if __name__ == "__main__":
    main()