server/app.py

"""OpenEnv-compatible server for the Crime Investigation Environment.

Exposes the CrimeInvestigationEnv via OpenEnv's HTTP/WebSocket interface.

Usage:
    # Development (with auto-reload):
    uvicorn server.app:app --reload --host 0.0.0.0 --port 8000

    # Or run directly:
    python -m server.app
"""

import os
import sys
import json
import base64
from typing import Any, Dict, List, Optional

from pydantic import Field

# Add project root to path so crime_env is importable
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from fastapi.responses import HTMLResponse

from openenv.core.env_server.http_server import create_app
from openenv.core.env_server.interfaces import Environment
from openenv.core.env_server.types import (
    Action,
    EnvironmentMetadata,
    Observation,
    State,
)

from crime_env.environment import CrimeInvestigationEnv


# ── Pydantic types for OpenEnv ──────────────────────────────────────────────


class CrimeAction(Action):
    """Action schema for the Crime Investigation environment."""

    action_string: str = Field(
        ...,
        description=(
            "Action in one of the following formats:\n"
            "  ACTION: ask_question | TARGET: <agent> | CONTENT: <question>\n"
            "  ACTION: request_evidence | ITEM: <item>\n"
            "  ACTION: accuse | TARGET: <suspect>"
        ),
        examples=[
            "ACTION: ask_question | TARGET: Suspect_A | CONTENT: Where were you?",
            "ACTION: request_evidence | ITEM: keycard_log",
            "ACTION: accuse | TARGET: Suspect_A",
        ],
    )


class CrimeObservation(Observation):
    """Observation schema returned by the Crime Investigation environment."""

    role: str = Field(default="detective", description="Agent role")
    briefing: str = Field(default="", description="Case briefing for the detective")
    turn: int = Field(default=0, description="Current turn number")
    conversation_history: List[Dict[str, Any]] = Field(
        default_factory=list, description="Full conversation history"
    )
    evidence_log: List[Dict[str, Any]] = Field(
        default_factory=list, description="Revealed evidence items"
    )
    message: str = Field(default="", description="System message for the current step")


class CrimeState(State):
    """State schema for the Crime Investigation environment."""

    turn: int = Field(default=0, description="Current turn number")
    is_done: bool = Field(default=False, description="Whether the episode is over")
    max_turns: int = Field(default=15, description="Maximum turns per episode")
    evidence_revealed: int = Field(default=0, description="Number of evidence items revealed")
    contradictions_found: int = Field(default=0, description="Number of contradictions detected")


# ── OpenEnv-compatible wrapper ──────────────────────────────────────────────


class CrimeInvestigationOpenEnv(Environment[CrimeAction, CrimeObservation, CrimeState]):
    """OpenEnv wrapper around CrimeInvestigationEnv."""

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._env = CrimeInvestigationEnv()
        self._current_obs: Optional[dict] = None

    def reset(
        self,
        seed: Optional[int] = None,
        episode_id: Optional[str] = None,
        **kwargs,
    ) -> CrimeObservation:
        if hasattr(self, "_reset_rubric"):
            self._reset_rubric()
        obs = self._env.reset()
        self._current_obs = obs
        return CrimeObservation(
            role=obs.get("role", "detective"),
            briefing=obs.get("briefing", ""),
            turn=obs.get("turn", 0),
            conversation_history=obs.get("conversation_history", []),
            evidence_log=obs.get("evidence_log", []),
            message=obs.get("message", ""),
            done=False,
            reward=None,
        )

    def step(
        self,
        action: CrimeAction,
        timeout_s: Optional[float] = None,
        **kwargs,
    ) -> CrimeObservation:
        obs_dict, reward, done, info = self._env.step(action.action_string)
        self._current_obs = obs_dict
        return CrimeObservation(
            role=obs_dict.get("role", "detective"),
            briefing=obs_dict.get("briefing", ""),
            turn=obs_dict.get("turn", 0),
            conversation_history=obs_dict.get("conversation_history", []),
            evidence_log=obs_dict.get("evidence_log", []),
            message=obs_dict.get("message", ""),
            done=done,
            reward=reward,
        )

    @property
    def state(self) -> CrimeState:
        env_state = self._env.state()
        return CrimeState(
            turn=env_state.get("turn", 0),
            is_done=env_state.get("done", False),
            max_turns=env_state.get("max_turns", 15),
            evidence_revealed=env_state.get("evidence_revealed", 0),
            contradictions_found=env_state.get("contradictions_found", 0),
        )

    def get_metadata(self) -> EnvironmentMetadata:
        return EnvironmentMetadata(
            name="CrimeInvestigationEnv",
            description=(
                "AI Crime Investigation World — a multi-agent RL environment "
                "where a detective agent interrogates suspects and a witness, "
                "reviews evidence, and makes an accusation."
            ),
            version="1.0.0",
        )

    def close(self) -> None:
        pass


# ── App creation ────────────────────────────────────────────────────────────

app = create_app(
    CrimeInvestigationOpenEnv,
    CrimeAction,
    CrimeObservation,
    env_name="crime_investigation",
    max_concurrent_envs=1,
)

# ── Custom Endpoints for Dashboard ──────────────────────────────────────────

@app.get("/", response_class=HTMLResponse)
async def serve_dashboard():
    """Serves the dashboard.html interface."""
    script_dir = os.path.dirname(os.path.abspath(__file__))
    project_root = os.path.dirname(script_dir)
    dashboard_path = os.path.join(project_root, "dashboard.html")
    with open(dashboard_path, "r", encoding="utf-8") as f:
        return f.read()

@app.get("/api/run_episode")
async def run_episode_api():
    """Runs a single test episode and returns the trace.
    
    Import is lazy (Issue 9) and execution is offloaded to a thread
    so the FastAPI event loop isn't blocked (Issue 11).
    """
    import asyncio
    from test_one_episode import run_test_episode

    rewards, info, trace = await asyncio.to_thread(run_test_episode)
    return {
        "status": "ok",
        "rewards": rewards,
        "info": info,
        "trace": trace
    }


def _moving_average(values: List[float], window: int) -> List[float]:
    if not values:
        return []
    if window <= 1:
        return values[:]
    averaged: List[float] = []
    running_sum = 0.0
    queue: List[float] = []
    for v in values:
        queue.append(float(v))
        running_sum += float(v)
        if len(queue) > window:
            running_sum -= queue.pop(0)
        averaged.append(running_sum / len(queue))
    return averaged


@app.get("/api/reward_curve")
async def reward_curve_api():
    """Return training reward history for dashboard/HF demo visibility."""
    script_dir = os.path.dirname(os.path.abspath(__file__))
    project_root = os.path.dirname(script_dir)
    rewards_path = os.path.join(project_root, "rewards.json")
    reward_curve_path = os.path.join(project_root, "reward_curve.png")

    rewards: List[float] = []
    results: List[str] = []
    model_name = "unknown"
    num_episodes = 0
    rewards_file_found = os.path.exists(rewards_path)

    if rewards_file_found:
        with open(rewards_path, "r", encoding="utf-8") as f:
            payload = json.load(f)
        rewards = [float(x) for x in payload.get("rewards", [])]
        results = [str(x) for x in payload.get("results", [])]
        model_name = str(payload.get("model", "unknown"))
        num_episodes = int(payload.get("num_episodes", len(rewards)))

    window = min(20, max(1, len(rewards) // 4))
    smoothed = _moving_average(rewards, window)
    mean_first = sum(rewards[:50]) / max(1, min(50, len(rewards))) if rewards else 0.0
    mean_last = sum(rewards[-50:]) / max(1, min(50, len(rewards))) if rewards else 0.0

    image_data_url = None
    if os.path.exists(reward_curve_path):
        with open(reward_curve_path, "rb") as f:
            encoded = base64.b64encode(f.read()).decode("ascii")
        image_data_url = f"data:image/png;base64,{encoded}"

    return {
        "status": "ok",
        "has_data": bool(rewards),
        "rewards_file_found": rewards_file_found,
        "message": (
            "Training data loaded"
            if rewards
            else "No rewards.json found on server yet. Commit and push training artifacts to update this panel."
        ),
        "model": model_name,
        "num_episodes": num_episodes,
        "rewards": rewards,
        "smoothed": smoothed,
        "smooth_window": window,
        "mean_first_50": round(mean_first, 4),
        "mean_last_50": round(mean_last, 4),
        "improvement": round(mean_last - mean_first, 4),
        "results": {
            "correct": results.count("correct"),
            "wrong": results.count("wrong"),
            "timeout": results.count("timeout"),
        },
        "image_data_url": image_data_url,
    }


@app.get("/api/health")
async def health_api():
    """Simple endpoint used for deployment sanity checks."""
    return {"status": "ok", "service": "crime-investigation"}


def main(host: str = "0.0.0.0", port: int = 8000):
    """Entry point for direct execution."""
    import uvicorn

    uvicorn.run(app, host=host, port=port)


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--port", type=int, default=8000)
    args = parser.parse_args()
    main(port=args.port)