| """ |
| FastAPI app for the BA Agent RL Environment (OpenEnv). |
| |
| Endpoints (provided by OpenEnv create_app): |
| - POST /reset, POST /step, GET /state, GET /schema, WS /ws |
| |
| Custom: |
| - GET /web — minimal demo UI (mirrors clinKriya base_path). |
| - GET /api/tasks — list available tasks (task_id, title, n_docs, n_golden). |
| - GET /api/tasks/{task_id} — fetch a task body. |
| """ |
| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
| from typing import Optional |
|
|
| from fastapi.responses import HTMLResponse, JSONResponse |
| from starlette.middleware.base import BaseHTTPMiddleware |
| from starlette.requests import Request |
|
|
| try: |
| from openenv.core.env_server.http_server import create_app |
| except Exception as exc: |
| raise ImportError( |
| "openenv-core required. Install with `uv pip install -e .` first." |
| ) from exc |
|
|
| from ba_agent_env.models import BAAgentAction, BAAgentObservation |
|
|
| from .ba_agent_environment import BAAgentEnvironment |
| from .dataset import get_tasks |
|
|
|
|
| _ui_env: Optional[BAAgentEnvironment] = None |
|
|
|
|
| def _get_ui_env() -> BAAgentEnvironment: |
| global _ui_env |
| if _ui_env is None: |
| _ui_env = BAAgentEnvironment() |
| return _ui_env |
|
|
|
|
| _ROOT = Path(__file__).resolve().parents[1] |
| _UI_HTML_PATH = _ROOT / "ui" / "index.html" |
|
|
|
|
| def _ui_html() -> str: |
| if _UI_HTML_PATH.exists(): |
| return _UI_HTML_PATH.read_text(encoding="utf-8") |
| tasks = get_tasks() |
| rows = "".join( |
| f"<tr><td>{t.task_id}</td><td>{t.title}</td><td>{len(t.input_documents)}</td><td>{len(t.golden_stories)}</td></tr>" |
| for t in tasks |
| ) |
| return f"""<!doctype html> |
| <html><head><meta charset='utf-8'><title>BA Agent RL Env</title> |
| <style>body{{font-family:system-ui;max-width:960px;margin:40px auto;padding:0 16px;color:#111}} |
| table{{width:100%;border-collapse:collapse;margin-top:16px}}th,td{{border:1px solid #ddd;padding:8px;text-align:left}} |
| code{{background:#f5f5f5;padding:2px 6px;border-radius:4px}} |
| .muted{{color:#666;font-size:14px}} |
| </style></head><body> |
| <h1>BA Agent RL Environment</h1> |
| <p>Enterprise requirements-generation RL environment. Agents run a 5-stage BA pipeline |
| (<code>EXTRACT</code> -> <code>INTERVIEW</code> -> <code>GRAPH</code> -> <code>STORY_GEN</code> -> <code>PRD</code> -> <code>FINISH</code>) over one feature |
| and are scored against BA-authored golden user stories.</p> |
| <p class='muted'>Endpoints: <code>POST /reset</code>, <code>POST /step</code>, <code>GET /state</code>, <code>GET /schema</code>, <code>GET /api/tasks</code>.</p> |
| <h2>Available tasks ({len(tasks)})</h2> |
| <table><thead><tr><th>task_id</th><th>title</th><th>docs</th><th>golden stories</th></tr></thead> |
| <tbody>{rows}</tbody></table> |
| <h2>Reward composition</h2> |
| <p>Alignment 35% · Coherence 24% · Completeness 18% · Compliance 10% · Testability 9% · Spec Quality 4%.</p> |
| <p class='muted'>LLM judge (OpenRouter / GPT-4o) when <code>OPENROUTER_API_KEY</code> is set; heuristic fallback otherwise.</p> |
| </body></html>""" |
|
|
|
|
| class _UIMiddleware(BaseHTTPMiddleware): |
| async def dispatch(self, request: Request, call_next): |
| p = request.url.path |
| if p in ("/", "/ui", "/web") or p.startswith("/web/"): |
| return HTMLResponse(content=_ui_html()) |
| return await call_next(request) |
|
|
|
|
| app = create_app( |
| BAAgentEnvironment, |
| BAAgentAction, |
| BAAgentObservation, |
| env_name="ba_agent_env", |
| max_concurrent_envs=1, |
| ) |
|
|
| app.add_middleware(_UIMiddleware) |
|
|
|
|
| @app.get("/api/tasks") |
| async def list_tasks(): |
| tasks = get_tasks() |
| return JSONResponse(content=[ |
| { |
| "task_id": t.task_id, |
| "title": t.title, |
| "n_input_documents": len(t.input_documents), |
| "n_golden_stories": len(t.golden_stories), |
| } |
| for t in tasks |
| ]) |
|
|
|
|
| @app.get("/api/tasks/{task_id}") |
| async def get_task(task_id: str): |
| from .dataset import get_task_by_id |
| t = get_task_by_id(task_id) |
| if not t: |
| return JSONResponse(content={"error": f"task '{task_id}' not found"}, status_code=404) |
| return JSONResponse(content=json.loads(t.model_dump_json())) |
|
|
|
|
| @app.get("/health") |
| async def health(): |
| return {"status": "ok", "n_tasks": len(get_tasks())} |
|
|
|
|
| @app.post("/api/reset") |
| async def api_reset(): |
| """Stateful reset for the UI — shared env instance across /api/step calls.""" |
| env = _get_ui_env() |
| obs = env.reset() |
| return JSONResponse(content={"observation": json.loads(obs.model_dump_json())}) |
|
|
|
|
| @app.post("/api/step") |
| async def api_step(request: Request): |
| """Stateful step for the UI — uses the shared env instance.""" |
| body = await request.json() |
| raw = body.get("action", body) |
| action = BAAgentAction(**raw) |
| env = _get_ui_env() |
| obs = env.step(action) |
| return JSONResponse(content={"observation": json.loads(obs.model_dump_json())}) |
|
|
