"""
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"| {t.task_id} | {t.title} | {len(t.input_documents)} | {len(t.golden_stories)} |
"
for t in tasks
)
return f"""
BA Agent RL Env
BA Agent RL Environment
Enterprise requirements-generation RL environment. Agents run a 5-stage BA pipeline
(EXTRACT -> INTERVIEW -> GRAPH -> STORY_GEN -> PRD -> FINISH) over one feature
and are scored against BA-authored golden user stories.
Endpoints: POST /reset, POST /step, GET /state, GET /schema, GET /api/tasks.
Available tasks ({len(tasks)})
| task_id | title | docs | golden stories |
|---|
{rows}
Reward composition
Alignment 35% · Coherence 24% · Completeness 18% · Compliance 10% · Testability 9% · Spec Quality 4%.
LLM judge (OpenRouter / GPT-4o) when OPENROUTER_API_KEY is set; heuristic fallback otherwise.
"""
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())})