Spaces:
Sleeping
Sleeping
| """ | |
| 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: # pragma: no cover | |
| 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) | |
| 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 | |
| ]) | |
| 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())) | |
| async def health(): | |
| return {"status": "ok", "n_tasks": len(get_tasks())} | |
| 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())}) | |
| 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())}) | |