Spaces:
Sleeping
Sleeping
| """ | |
| Minimal HTTP client for the BA Agent RL Environment. | |
| Talks to either: | |
| - a local container (http://localhost:8000) | |
| - a HuggingFace Space (https://<user>-<space>.hf.space) | |
| Usage: | |
| from client import BAAgentClient | |
| env = BAAgentClient("http://localhost:8000") | |
| obs = env.reset() | |
| while not done: | |
| obs, reward, done, info = env.step(action_type="EXTRACT", payload="...") | |
| """ | |
| from __future__ import annotations | |
| import json | |
| from typing import Any, Dict, List, Optional, Tuple | |
| import requests | |
| class BAAgentClient: | |
| def __init__(self, base_url: str = "http://localhost:8000", timeout: int = 30): | |
| self.base_url = base_url.rstrip("/") | |
| self.timeout = timeout | |
| def reset(self) -> Dict[str, Any]: | |
| r = requests.post(f"{self.base_url}/reset", json={}, timeout=self.timeout) | |
| r.raise_for_status() | |
| return (r.json() or {}).get("observation", r.json()) | |
| def step(self, action_type: str, payload: str = "") -> Tuple[Dict[str, Any], float, bool, Dict[str, Any]]: | |
| body = {"action": {"action_type": action_type, "payload": payload}} | |
| r = requests.post(f"{self.base_url}/step", json=body, timeout=self.timeout) | |
| r.raise_for_status() | |
| j = r.json() | |
| obs = j.get("observation", j) | |
| reward = float(j.get("reward", 0.0)) | |
| done = bool(j.get("done", False)) | |
| meta = j.get("metadata", {}) | |
| return obs, reward, done, meta | |
| def state(self) -> Dict[str, Any]: | |
| r = requests.get(f"{self.base_url}/state", timeout=self.timeout) | |
| r.raise_for_status() | |
| return r.json() | |
| def tasks(self) -> List[Dict[str, Any]]: | |
| r = requests.get(f"{self.base_url}/api/tasks", timeout=self.timeout) | |
| r.raise_for_status() | |
| return r.json() | |
| def task(self, task_id: str) -> Dict[str, Any]: | |
| r = requests.get(f"{self.base_url}/api/tasks/{task_id}", timeout=self.timeout) | |
| r.raise_for_status() | |
| return r.json() | |
| if __name__ == "__main__": | |
| env = BAAgentClient() | |
| obs = env.reset() | |
| print(f"Task {obs['task_id']} :: {obs['title']}") | |
| stub_stories = [ | |
| {"title": "Receive document", "description": "As an operator, I want to receive an inbound document, so that I can begin processing.", "acceptance_criteria": "Given inbound, When parsed, Then Document record created."}, | |
| {"title": "Validate document", "description": "As an operator, I want validation, so that bad data is rejected early.", "acceptance_criteria": "Given Document, When validated, Then errors flagged."}, | |
| {"title": "Finalise document", "description": "As an operator, I want to finalise, so that the record is auditable.", "acceptance_criteria": "Given validated Document, When finalised, Then status=Finalised."}, | |
| ] | |
| plan = [ | |
| ("EXTRACT", "primary workflow, entities=[Actor, System, Document], constraints=compliance."), | |
| ("INTERVIEW", json.dumps([{"q": "Primary actor?", "a": "Operator"}])), | |
| ("GRAPH", json.dumps({"nodes": ["Actor", "Document"], "edges": [["Actor", "Document"]]})), | |
| ("STORY_GEN", json.dumps(stub_stories)), | |
| ("FINISH", ""), | |
| ] | |
| total = 0.0 | |
| for at, pl in plan: | |
| obs, r, done, meta = env.step(at, pl) | |
| total += r | |
| print(f"{at:>10} -> r={r:+.3f} done={done} {meta.get('status','')}") | |
| print(f"\nEpisode reward = {total:.3f}") | |
| if "composite_0_to_100" in meta: | |
| print(f"Composite (bench) = {meta['composite_0_to_100']}/100 ({meta.get('engine')})") | |