Spaces:

Addyk24
/

Project-Polymath

Sleeping

File size: 14,692 Bytes

6e2f008

"""
FastAPI server exposing the Project Polymath Workspace Environment.
Endpoints:
    GET  /              — Command Center UI (HTML)
    GET  /docs          — Interactive OpenAPI (Swagger UI)
    GET  /health        — Liveness probe (JSON)
    POST /reset         — Start a new negotiation
    POST /step          — Apply an agent action
    GET  /state         — Read current state without advancing
"""

from __future__ import annotations
import os
import sys

# Ensure Python can find your 'envs' and 'models' folders
ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
if ROOT_DIR not in sys.path:
    sys.path.insert(0, ROOT_DIR)

from typing import Any, Optional
from fastapi import FastAPI
from fastapi.responses import HTMLResponse
from pydantic import BaseModel, ConfigDict, Field
import uvicorn

# Project Polymath Imports
from envs.environment import WorkSpaceEnvironment 
from models.schemas import WorkSpaceAction

# ── OpenAPI Documentation Setup ──────────────────────────────────────────────

_OPENAPI_TAGS = [
    {
        "name": "Environment",
        "description": "Episode lifecycle for Polymath: call **reset** before **step**.",
    },
    {
        "name": "Interface",
        "description": "Browser UI for manual debugging (HTML, not JSON).",
    }
]

app = FastAPI(
    title="Project Polymath: OpenEnv",
    version="1.0.0",
    openapi_tags=_OPENAPI_TAGS,
    description=(
        "Multi-agent negotiation environment. "
        "Agent must balance constraints from Finance, Security, and UX."
    ),
)

# Force the environment into the mode you want the judges to see by default
os.environ["BASELINE_ENV_MODE"] = "easy"
_env = WorkSpaceEnvironment()

# ── request / response schemas ───────────────────────────────────────────────

class ResetRequest(BaseModel):
    """Start a new episode."""
    model_config = ConfigDict(
        json_schema_extra={
            "examples": [{"topic": "Draft the new Mobile App PRD"}]
        }
    )
    topic: str = Field(
        default="Draft the new Mobile App PRD",
        description="The core task the PM must complete."
    )

class StepRequest(BaseModel):
    """Apply one agent action."""
    action: WorkSpaceAction = Field(
        description="The agent's action payload."
    )

class EnvResponse(BaseModel):
    """Observation and reward after reset, step, or state."""
    observation: dict[str, Any] | None = Field(description="Environment feedback and turn count.")
    reward: float = Field(description="Reward for the last transition.")
    done: bool = Field(description="True after the episode concludes.")
    info: dict[str, Any] = Field(default_factory=dict)

def _format_obs(obs: Any) -> dict[str, Any] | None:
    if obs is None:
        return None
    # obs is a WorkspaceObservation
    return {
        "feedback": getattr(obs, "feedback", "Episode Terminated."),
        "current_turn": getattr(obs, "current_turn", 0)
    }

# ── FRONTEND HTML/CSS/JS PAYLOAD ─────────────────────────────────────────────

_DEBUG_UI_HTML = """<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1" />
  <title>Polymath Command Center</title>
  <style>
    :root {
      --bg: #0f172a; --surface: #1e293b; --border: #334155; --text: #f8fafc;
      --muted: #94a3b8; --accent: #8b5cf6; --exec: #8b5cf6; --exec-hover: #7c3aed;
      --card-blue: #3b82f6; --card-orange: #f97316; --card-green: #10b981; --card-grey: #475569;
      --mono: ui-monospace, "Cascadia Code", monospace;
    }
    * { box-sizing: border-box; }
    body { margin: 0; min-height: 100vh; font-family: system-ui, sans-serif; background: var(--bg); color: var(--text); }
    .wrap { max-width: 60rem; margin: 0 auto; padding: 2rem 1rem; }
    
    .page-head { display: flex; justify-content: space-between; align-items: flex-start; margin-bottom: 1.5rem; }
    h1 { font-size: 1.5rem; margin: 0 0 0.25rem; }
    .sub { color: var(--muted); font-size: 0.9rem; margin: 0; }
    .btn-mini { background: #fdfd96; color: #111; padding: 0.35rem 0.65rem; border-radius: 6px; text-decoration: none; font-weight: bold; font-size: 0.75rem; text-transform: uppercase; }

    .panel { background: var(--surface); border: 1px solid var(--border); border-radius: 8px; padding: 1.25rem; margin-bottom: 1.25rem; }
    .metrics { display: grid; grid-template-columns: repeat(4, 1fr); gap: 1rem; margin-bottom: 1.5rem; }
    .metric { padding: 1rem; border-radius: 8px; color: #fff; box-shadow: 0 4px 6px rgba(0,0,0,0.3); border-top: 4px solid; background: #1e293b; }
    .metric .m-label { font-size: 0.7rem; text-transform: uppercase; color: var(--muted); font-weight: bold; }
    .metric .m-val { font-size: 1.5rem; font-weight: bold; margin-top: 0.25rem; }
    
    label { font-size: 0.8rem; color: var(--muted); font-weight: bold; display: block; margin-bottom: 0.4rem; text-transform: uppercase; }
    select, input, textarea { width: 100%; padding: 0.75rem; border-radius: 6px; border: 1px solid var(--border); background: #0f172a; color: var(--text); font-family: inherit; margin-bottom: 1rem;}
    textarea { min-height: 80px; resize: vertical; }
    
    .btn-row { display: flex; gap: 0.75rem; flex-wrap: wrap; margin-top: 0.5rem; }
    button { padding: 0.75rem 1.25rem; border: none; border-radius: 6px; font-weight: bold; cursor: pointer; transition: 0.2s; text-transform: uppercase; font-size: 0.8rem; }
    .btn-exec { background: var(--exec); color: #fff; flex: 2; }
    .btn-exec:hover { background: var(--exec-hover); }
    .btn-sec { background: var(--card-grey); color: #fff; flex: 1; }
    .btn-sec:hover { background: #334155; }
    
    .world-text { background: #0f172a; padding: 1rem; border-radius: 6px; border: 1px solid var(--border); line-height: 1.5; color: #e2e8f0; white-space: pre-wrap; font-family: var(--mono); font-size: 0.85rem;}
    
    .timeline { max-height: 400px; overflow-y: auto; display: flex; flex-direction: column; gap: 0.75rem; padding-right: 0.5rem; }
    .tl-item { background: #0f172a; border: 1px solid var(--border); padding: 1rem; border-radius: 6px; }
    .tl-header { display: flex; justify-content: space-between; margin-bottom: 0.5rem; border-bottom: 1px solid var(--border); padding-bottom: 0.5rem;}
    .tl-reward { font-weight: bold; font-family: var(--mono); }
    .tl-pos { color: #10b981; } .tl-neg { color: #ef4444; }
    .tl-body { font-size: 0.85rem; line-height: 1.4; color: var(--muted); }
    .tl-body strong { color: #e2e8f0; }

    #toast { position: fixed; top: 1rem; right: 1rem; background: #ef4444; color: white; padding: 1rem; border-radius: 6px; display: none; z-index: 1000; box-shadow: 0 4px 12px rgba(0,0,0,0.5); }
  </style>
</head>
<body>
  <div id="toast"></div>
  <div class="wrap">
    
    <header class="page-head">
      <div>
        <h1>Project Polymath</h1>
        <p class="sub">Multi-Agent Negotiation Environment • <a href="/health">/health</a></p>
      </div>
      <a class="btn-mini" href="/docs" target="_blank">Docs</a>
    </header>
    
    <div class="metrics">
      <div class="metric" style="border-color: var(--card-blue);">
        <div class="m-label">Step Reward</div><div class="m-val" id="mStep">0.00</div>
      </div>
      <div class="metric" style="border-color: var(--card-orange);">
        <div class="m-label">Total Reward</div><div class="m-val" id="mTotal">0.00</div>
      </div>
      <div class="metric" style="border-color: var(--card-green);">
        <div class="m-label">Turn Count</div><div class="m-val" id="mTurn">0</div>
      </div>
      <div class="metric" style="border-color: var(--card-grey);">
        <div class="m-label">Status</div><div class="m-val" id="mStatus" style="font-size: 1.1rem; padding-top: 0.4rem;">STANDBY</div>
      </div>
    </div>

    <div class="panel">
      <label>Environment Feedback</label>
      <div class="world-text" id="oppText">Initialize the environment to begin.</div>
    </div>

    <div class="panel">
      <label>Topic / Seed</label>
      <input type="text" id="iTopic" value="Draft the new Mobile App PRD">
      
      <div style="display: flex; gap: 1rem;">
        <div style="flex: 1;">
          <label>Action Type</label>
          <select id="iType" onchange="toggleTarget()">
            <option value="message_expert">message_expert</option>
            <option value="propose_draft">propose_draft</option>
            <option value="submit_final">submit_final</option>
          </select>
        </div>
        <div style="flex: 1;">
          <label>Target</label>
          <select id="iTarget">
            <option value="Finance">Finance</option>
            <option value="Security">Security</option>
            <option value="UX">UX</option>
            <option value="All">All</option>
          </select>
        </div>
      </div>
      
      <label>Content Payload</label>
      <textarea id="iArg" placeholder="Message the expert or paste the PRD..."></textarea>
      
      <div class="btn-row">
        <button class="btn-exec" onclick="doStep()">Execute Step</button>
        <button class="btn-sec" onclick="doState()">Refresh State</button>
        <button class="btn-sec" onclick="doReset()">Reset Env</button>
      </div>
    </div>

    <div class="panel">
      <label>Action Timeline</label>
      <div class="timeline" id="timeline">
        <div style="color: var(--muted); text-align: center; padding: 2rem;">No actions yet. Click Reset to begin.</div>
      </div>
    </div>

  </div>

<script>
  let totalReward = 0;
  let isDone = false;

  const $ = id => document.getElementById(id);
  
  function showToast(msg) {
    const t = $('toast'); t.innerText = msg; t.style.display = 'block';
    setTimeout(() => t.style.display = 'none', 4000);
  }

  function toggleTarget() {
    const type = $('iType').value;
    $('iTarget').disabled = (type === 'submit_final');
  }

  function updateUI(data, isReset = false) {
    if(isReset) { totalReward = 0; $('timeline').innerHTML = ''; }
    
    const obs = data.observation || {};
    const reward = data.reward || 0;
    totalReward += reward;
    isDone = data.done;
    
    $('mStep').innerText = reward.toFixed(2);
    $('mTotal').innerText = totalReward.toFixed(2);
    $('mTurn').innerText = obs.current_turn || 0;
    $('mStatus').innerText = data.done ? "TERMINATED" : "ACTIVE";
    $('oppText').innerText = obs.feedback || "Episode Terminated.";

    if(data.action_taken) {
      const colorCls = reward >= 0 ? "tl-pos" : "tl-neg";
      const sign = reward > 0 ? "+" : "";
      const html = `
        <div class="tl-item">
          <div class="tl-header">
            <strong style="color: #fff;">${data.action_taken.action_type} &rarr; ${data.action_taken.target || "None"}</strong>
            <span class="tl-reward ${colorCls}">${sign}${reward.toFixed(2)}</span>
          </div>
          <div class="tl-body">
            <div><strong>Agent:</strong> ${data.action_taken.content}</div>
          </div>
        </div>`;
      $('timeline').insertAdjacentHTML('beforeend', html);
      $('timeline').scrollTop = $('timeline').scrollHeight; 
    }
  }

  async function apiCall(endpoint, payload) {
    try {
      const res = await fetch(endpoint, {
        method: payload ? "POST" : "GET",
        headers: { "Content-Type": "application/json" },
        body: payload ? JSON.stringify(payload) : null
      });
      const data = await res.json();
      if(!res.ok) throw new Error(data.detail || "API Error");
      return data;
    } catch (e) { showToast(e.message); throw e; }
  }

  async function doReset() {
    const data = await apiCall("/reset", { topic: $('iTopic').value });
    updateUI(data, true);
    $('iArg').value = ""; 
  }

  async function doState() {
    const data = await apiCall("/state");
    updateUI(data);
  }

  async function doStep() {
    if(isDone || $('mStatus').innerText === "STANDBY") {
      showToast("Please Reset the environment first!"); return;
    }
    
    let targetVal = $('iTarget').value;
    if ($('iType').value === "submit_final") targetVal = null;

    const action = { 
      action_type: $('iType').value, 
      target: targetVal,
      content: $('iArg').value 
    };

    const data = await apiCall("/step", { action });
    data.action_taken = action; 
    updateUI(data);
    $('iArg').value = ""; 
  }
</script>
</body>
</html>
"""

# ── routes ───────────────────────────────────────────────────────────────────

@app.get("/", response_class=HTMLResponse, tags=["Interface"])
def ui() -> HTMLResponse:
    """Browser debug UI (HTML)."""
    return HTMLResponse(content=_DEBUG_UI_HTML)

@app.get("/health", tags=["System"])
async def health() -> dict[str, str]:
    """Return `{"status": "ok"}` when the server is up."""
    return {"status": "ok"}

@app.post("/reset", response_model=EnvResponse, tags=["Environment"])
async def reset(req: ResetRequest | None = None) -> dict[str, Any]:
    if req is None:
        req = ResetRequest()
    obs = _env.reset(req.topic)
    return {
        "observation": _format_obs(obs),
        "reward": 0.0,
        "done": False,
        "info": {}
    }

@app.post("/step", response_model=EnvResponse, tags=["Environment"])
async def step(req: StepRequest) -> dict[str, Any]:
    obs = _env.step(req.action)
    return {
        "observation": _format_obs(obs),
        "reward": float(obs.reward),
        "done": obs.done,
        "info": {}
    }

@app.get("/state", response_model=EnvResponse, tags=["Environment"])
async def state() -> dict[str, Any]:
    try:
        obs = _env._state # Direct access for API state check
        if obs is None:
             return {"observation": None, "reward": 0.0, "done": False, "info": {}}
        return {
            "observation": {
                "feedback": "Current state fetched.",
                "current_turn": obs.turn_count
            },
            "reward": 0.0,
            "done": obs.is_done,
            "info": {}
        }
    except Exception:
        return {"observation": None, "reward": 0.0, "done": False, "info": {}}

def main() -> None:
    uvicorn.run(app, host="0.0.0.0", port=7860)

if __name__ == "__main__":
    main()