|
| 1 | +"""Client library for MissionCtrl environment. |
| 2 | +
|
| 3 | +Provides HTTP client functions for interacting with the MissionCtrl environment API. |
| 4 | +Contains environment payload structures and example scenarios. |
| 5 | +""" |
| 6 | + |
| 7 | +import os |
| 8 | +from typing import Any, Dict, Optional |
| 9 | + |
| 10 | +import httpx |
| 11 | +from dotenv import load_dotenv |
| 12 | + |
| 13 | +# --------------------------------------------------------------------------- |
| 14 | +# Load .env file automatically |
| 15 | +# --------------------------------------------------------------------------- |
| 16 | +load_dotenv() |
| 17 | + |
| 18 | +# --------------------------------------------------------------------------- |
| 19 | +# Environment configuration |
| 20 | +# --------------------------------------------------------------------------- |
| 21 | +ENV_BASE_URL: str = os.environ.get("ENV_BASE_URL", "http://localhost:7860") |
| 22 | + |
| 23 | +# Known agent types in the environment |
| 24 | +KNOWN_AGENTS = ( |
| 25 | + "PlannerAgent", |
| 26 | + "ResearchAgent", |
| 27 | + "CoderAgent", |
| 28 | + "TesterAgent", |
| 29 | + "CommAgent", |
| 30 | +) |
| 31 | + |
| 32 | +# Available task tiers |
| 33 | +TASKS = ["easy", "medium", "hard", "special"] |
| 34 | + |
| 35 | +# Default max steps per episode |
| 36 | +MAX_STEPS = 5 |
| 37 | + |
| 38 | + |
| 39 | +# --------------------------------------------------------------------------- |
| 40 | +# HTTP Client |
| 41 | +# --------------------------------------------------------------------------- |
| 42 | +http = httpx.Client(timeout=60.0) |
| 43 | + |
| 44 | + |
| 45 | +# --------------------------------------------------------------------------- |
| 46 | +# Environment API Functions |
| 47 | +# --------------------------------------------------------------------------- |
| 48 | +def reset_env(task_id: str, seed: Optional[int] = None) -> Dict[str, Any]: |
| 49 | + """Reset the environment for a specific task. |
| 50 | + |
| 51 | + Args: |
| 52 | + task_id: The task tier to run (easy, medium, hard, special) |
| 53 | + seed: Optional random seed for reproducibility |
| 54 | + |
| 55 | + Returns: |
| 56 | + Dictionary containing the initial observation |
| 57 | + """ |
| 58 | + payload = {"task_id": task_id} |
| 59 | + if seed is not None: |
| 60 | + payload["seed"] = seed |
| 61 | + |
| 62 | + resp = http.post(f"{ENV_BASE_URL}/reset", json=payload) |
| 63 | + resp.raise_for_status() |
| 64 | + return resp.json() |
| 65 | + |
| 66 | + |
| 67 | +def step_env(action: str) -> Dict[str, Any]: |
| 68 | + """Execute one action in the environment. |
| 69 | + |
| 70 | + Args: |
| 71 | + action: The action string to execute (e.g., "APPROVE(task_1)") |
| 72 | + |
| 73 | + Returns: |
| 74 | + Dictionary containing the new observation, reward, done flag, and info |
| 75 | + """ |
| 76 | + resp = http.post(f"{ENV_BASE_URL}/step", json={"action": action}) |
| 77 | + resp.raise_for_status() |
| 78 | + return resp.json() |
| 79 | + |
| 80 | + |
| 81 | +def get_state() -> Dict[str, Any]: |
| 82 | + """Get the current environment state (read-only). |
| 83 | + |
| 84 | + Returns: |
| 85 | + Dictionary containing the current observation |
| 86 | + """ |
| 87 | + resp = http.get(f"{ENV_BASE_URL}/state") |
| 88 | + resp.raise_for_status() |
| 89 | + return resp.json() |
| 90 | + |
| 91 | + |
| 92 | +def get_history() -> list: |
| 93 | + """Get the action history for the current episode. |
| 94 | + |
| 95 | + Returns: |
| 96 | + List of past actions and their results |
| 97 | + """ |
| 98 | + resp = http.get(f"{ENV_BASE_URL}/history") |
| 99 | + resp.raise_for_status() |
| 100 | + return resp.json() |
| 101 | + |
| 102 | + |
| 103 | +def record_result(tier: str, score: float, steps: int, history: list, |
| 104 | + score_breakdown: Optional[Dict] = None, |
| 105 | + hallucination_stats: Optional[Dict] = None) -> Dict[str, str]: |
| 106 | + """Push a completed episode result to the dashboard. |
| 107 | + |
| 108 | + Args: |
| 109 | + tier: Task tier (easy, medium, hard, special) |
| 110 | + score: Final score for the episode |
| 111 | + steps: Number of steps taken |
| 112 | + history: Action history for the episode |
| 113 | + score_breakdown: Optional detailed score breakdown |
| 114 | + hallucination_stats: Optional hallucination detection statistics |
| 115 | + |
| 116 | + Returns: |
| 117 | + Confirmation response |
| 118 | + """ |
| 119 | + payload = { |
| 120 | + "tier": tier, |
| 121 | + "score": score, |
| 122 | + "steps": steps, |
| 123 | + "history": history, |
| 124 | + "score_breakdown": score_breakdown or {}, |
| 125 | + "hallucination_stats": hallucination_stats or {}, |
| 126 | + } |
| 127 | + resp = http.post(f"{ENV_BASE_URL}/record", json=payload) |
| 128 | + resp.raise_for_status() |
| 129 | + return resp.json() |
| 130 | + |
| 131 | + |
| 132 | +# --------------------------------------------------------------------------- |
| 133 | +# Example Usage |
| 134 | +# --------------------------------------------------------------------------- |
| 135 | +def example_basic_usage(): |
| 136 | + """Example showing basic environment interaction.""" |
| 137 | + # Reset environment for easy task |
| 138 | + result = reset_env("easy") |
| 139 | + obs = result["observation"] |
| 140 | + print(f"Started task with {len(obs['tasks'])} tasks") |
| 141 | + |
| 142 | + # Take a step |
| 143 | + step_result = step_env("NOOP") |
| 144 | + print(f"Reward: {step_result['reward']}, Done: {step_result['done']}") |
| 145 | + |
| 146 | + # Get current state |
| 147 | + state = get_state() |
| 148 | + print(f"Current step: {state['time_step']}") |
| 149 | + |
| 150 | + |
| 151 | +if __name__ == "__main__": |
| 152 | + # For backward compatibility, delegate to inference script |
| 153 | + from inference import main |
| 154 | + main() |
0 commit comments