""" OpenGrid Inference Script ========================= Runs an LLM agent against all OpenGrid tasks via the OpenAI-compatible API. Supports both single-agent and multi-agent POMDP modes. Optional environment variables: API_BASE_URL -- defaults to https://api.openai.com/v1 MODEL_NAME -- defaults to gpt-4o Required (one of): OPENAI_API_KEY or HF_TOKEN Emits structured [START], [STEP], [END] logs to stdout. Usage: # Single-agent mode (backward compatible) python inference.py # Multi-agent mode (uses safety layer + oversight) python inference.py --multi """ import os import sys import json import math import argparse import httpx from openai import OpenAI # ---------- Configuration ---------- API_BASE_URL = os.environ.get("API_BASE_URL", "https://api.openai.com/v1") MODEL_NAME = os.environ.get("MODEL_NAME", "gpt-4o") # Prefer OPENAI_API_KEY when using OpenAI endpoint; otherwise try HF_TOKEN API_KEY = ( os.environ.get("API_KEY") or os.environ.get("OPENAI_API_KEY") or os.environ.get("HF_TOKEN") or "" ) ENV_URL = os.environ.get("ENV_URL", "http://localhost:7860") BENCHMARK = "OpenGrid" # Safety cap — the environment's 'done' flag is the true terminator MAX_STEPS = 100 SUCCESS_SCORE_THRESHOLD = 0.5 TASKS = ["task_easy", "task_medium", "task_hard", "karnataka_easy", "karnataka_medium", "karnataka_hard", "task_karnataka"] SYSTEM_PROMPT_SINGLE = """You are a Power Grid Controller AI. Your goal is to maintain grid stability. Key objectives: 1. Keep grid frequency close to 50.0 Hz (acceptable: 49.5-50.5 Hz) 2. Prevent transmission line overloads (rho < 1.0) 3. Avoid grid islanding (blackout) Available actions: 1. bus_adjustments: List of {"bus_id": int, "delta": float} - Positive delta = increase power injection (discharge battery / ramp up generator) - Negative delta = decrease power injection (charge battery / ramp down generator) - Only works on battery and generator buses (avoid adjusting the slack bus — physics overwrites it) 2. topology_actions: List of {"line_id": str, "action": "open" | "close"} - Opening a line removes it; closing reconnects. 3-step cooldown. - WARNING: Opening lines can cause islanding -> blackout - Prefer NO topology actions unless absolutely necessary. Always return "topology_actions": [] Strategy: - If frequency < 50 Hz -> discharge batteries, ramp up generators - If frequency > 50 Hz -> charge batteries, ramp down generators - If a line rho > 0.9 -> reduce generation near that line, do NOT open it - Prefer minimal actions over aggressive switching Respond with ONLY a valid JSON object. Example: {"bus_adjustments": [{"bus_id": 2, "delta": 5.0}], "topology_actions": []} """ SYSTEM_PROMPT_MULTI = """You are a KPTCL Zone Controller AI managing one zone of the Karnataka power grid. You can only see and control buses in YOUR zone. Other zones are managed by other agents. Key objectives: 1. Keep grid frequency close to 50.0 Hz (you see a noisy reading) 2. Prevent line overloads in your zone (rho < 1.0) 3. Coordinate with other zones (don't fight against them) 4. Avoid actions that would trigger the safety layer Available actions: 1. bus_adjustments: List of {"bus_id": int, "delta": float} - ONLY adjust battery and generator buses in YOUR zone (avoid slack — physics overwrites it) - Positive delta = increase power injection - Negative delta = decrease power injection 2. topology_actions: List of {"line_id": str, "action": "open" | "close"} - Only for visible internal or boundary lines. Safety layer will block dangerous switches. - Boundary-line switching is risky; avoid unless necessary. Strategy: - If frequency < 50 Hz -> increase generation/discharge in your zone - If frequency > 50 Hz -> decrease generation/charge in your zone - Check neighbor signals to understand if other zones are compensating - Prefer small corrections over large swings Respond with ONLY a valid JSON object. Example: {"bus_adjustments": [{"bus_id": 2, "delta": 5.0}], "topology_actions": []} """ # ---------- Structured Logging ---------- def log_start(task: str, env: str, model: str, mode: str = "single"): print(f"[START] task={task} env={env} model={model} mode={mode}", flush=True) def log_step(step: int, action: str, reward: float, done: bool, error=None, agent_id=None): done_val = str(done).lower() error_val = str(error) if error else "null" agent_str = f" agent={agent_id}" if agent_id is not None else "" print( f"[STEP] step={step}{agent_str} action={action} reward={reward:.2f} done={done_val} error={error_val}", flush=True, ) def clamp_score(s: float) -> float: """Ensure score is strictly in (0, 1). Mirrors grader._clamp_score.""" try: s = float(s) except (TypeError, ValueError): return 0.5 if not math.isfinite(s): return 0.5 s = max(0.02, min(0.98, s)) s = math.floor(s * 10000) / 10000 return max(0.02, min(0.98, s)) def log_end(success: bool, steps: int, score: float, rewards: list, mode: str = "single"): clamped = clamp_score(score) success_val = str(success).lower() rewards_str = ",".join(f"{r:.2f}" for r in rewards) print( f"[END] success={success_val} steps={steps} score={clamped:.4f} rewards={rewards_str} mode={mode}", flush=True, ) # ---------- LLM Call ---------- def get_model_message(client: OpenAI, step: int, obs_json: str, last_reward: float, history: list, system_prompt: str, zone_name: str = None) -> str: """Ask the LLM what action to take given the current observation.""" context = "" if zone_name: context += f"[Zone: {zone_name}] " context += f"Step {step} | Last reward: {last_reward:+.2f}\n" if history: context += "Recent history (last 3):\n" + "\n".join(history[-3:]) + "\n\n" context += f"Current Grid State:\n{obs_json}" try: kwargs = dict( model=MODEL_NAME, messages=[ {"role": "system", "content": system_prompt}, {"role": "user", "content": context} ], temperature=0.0, max_tokens=300, ) # Use JSON mode if the endpoint supports it (OpenAI-compatible) try: kwargs["response_format"] = {"type": "json_object"} response = client.chat.completions.create(**kwargs) except Exception: # Fallback: endpoint may not support response_format kwargs.pop("response_format", None) response = client.chat.completions.create(**kwargs) return response.choices[0].message.content.strip() except Exception as exc: print(f"[DEBUG] Model request failed: {exc}", flush=True) return '{"bus_adjustments": [], "topology_actions": []}' # ---------- Environment Client ---------- class EnvClient: """HTTP client for the OpenGrid FastAPI environment.""" def __init__(self, base_url: str): self.base_url = base_url.rstrip("/") self.client = httpx.Client( timeout=httpx.Timeout(connect=10.0, read=60.0, write=30.0, pool=10.0) ) self.session_id = None # --- Single-Agent --- def reset(self, task_id: str) -> dict: resp = self.client.post(f"{self.base_url}/reset", params={"task_id": task_id}) resp.raise_for_status() data = resp.json() self.session_id = data["session_id"] return data["observation"] def step(self, action_dict: dict) -> dict: resp = self.client.post( f"{self.base_url}/step", params={"session_id": self.session_id}, json=action_dict ) resp.raise_for_status() return resp.json() # --- Multi-Agent --- def reset_multi(self, task_id: str) -> dict: resp = self.client.post(f"{self.base_url}/reset_multi", params={"task_id": task_id}) resp.raise_for_status() data = resp.json() self.session_id = data["session_id"] return data def step_multi(self, agent_actions: dict) -> dict: resp = self.client.post( f"{self.base_url}/step_multi", params={"session_id": self.session_id}, json={"agent_actions": agent_actions} ) resp.raise_for_status() return resp.json() # --- Shared --- def state(self) -> dict: resp = self.client.get(f"{self.base_url}/state", params={"session_id": self.session_id}) resp.raise_for_status() return resp.json() def grade(self) -> dict: # Grading can trigger lazy bound estimation (multiple rollouts) — use long timeout resp = self.client.get( f"{self.base_url}/grader", params={"session_id": self.session_id}, timeout=180.0, ) resp.raise_for_status() return resp.json() def list_tasks(self) -> list: """Fetch available tasks from the server.""" resp = self.client.get(f"{self.base_url}/tasks") resp.raise_for_status() return resp.json() def close(self): self.client.close() # ---------- Parse Action ---------- NOOP_ACTION = {"bus_adjustments": [], "topology_actions": []} def parse_action(response_text: str) -> dict: """Parse LLM JSON response into an action dict. Handles markdown fences, prose preambles, JSON lists, and malformed output. """ try: text = str(response_text).strip() # Strip markdown code fences if text.startswith("```"): lines = text.splitlines() if lines and lines[0].startswith("```"): lines = lines[1:] if lines and lines[-1].startswith("```"): lines = lines[:-1] text = "\n".join(lines).strip() # Extract first JSON object from any surrounding prose start = text.find("{") end = text.rfind("}") if start < 0 or end <= start: return dict(NOOP_ACTION) data = json.loads(text[start:end + 1]) # Handle list wrapping (e.g. [{...}]) if isinstance(data, list): data = data[0] if data else {} if not isinstance(data, dict): return dict(NOOP_ACTION) bus_adjustments = data.get("bus_adjustments", []) topology_actions = data.get("topology_actions", []) if not isinstance(bus_adjustments, list): bus_adjustments = [] if not isinstance(topology_actions, list): topology_actions = [] return { "bus_adjustments": bus_adjustments, "topology_actions": topology_actions, } except Exception: return dict(NOOP_ACTION) # ---------- Single-Agent Runner ---------- def run_task_single(client: OpenAI, env: EnvClient, task_id: str) -> dict: """Run one task in single-agent mode and return results.""" history_msgs = [] rewards = [] steps_taken = 0 score = 0.05 success = False log_start(task=task_id, env=BENCHMARK, model=MODEL_NAME, mode="single") try: obs = env.reset(task_id) last_reward = 0.0 for step_num in range(1, MAX_STEPS + 1): obs_json = json.dumps(obs, indent=2) message = get_model_message(client, step_num, obs_json, last_reward, history_msgs, SYSTEM_PROMPT_SINGLE) action_dict = parse_action(message) result = env.step(action_dict) obs = result["observation"] reward = result.get("reward", {}).get("value", 0.0) done = result.get("done", False) rewards.append(reward) steps_taken = step_num last_reward = reward action_summary = json.dumps(action_dict) if len(action_summary) > 200: action_summary = action_summary[:200] + "..." log_step(step=step_num, action=action_summary, reward=reward, done=done) history_msgs.append(f"Step {step_num}: action={action_summary[:80]} -> reward {reward:+.2f}") if done: break grade_result = env.grade() score = clamp_score(grade_result.get("score", 0.5)) success = score >= SUCCESS_SCORE_THRESHOLD except Exception as e: print(f"[DEBUG] Task {task_id} error: {e}", flush=True) score = 0.05 success = False log_end(success=success, steps=steps_taken, score=score, rewards=rewards, mode="single") return {"task": task_id, "score": score, "steps": steps_taken, "success": success} # ---------- Multi-Agent Runner ---------- def run_task_multi(client: OpenAI, env: EnvClient, task_id: str) -> dict: """Run one task in multi-agent mode and return results.""" rewards = [] steps_taken = 0 score = 0.05 success = False total_safety_interventions = 0 log_start(task=task_id, env=BENCHMARK, model=MODEL_NAME, mode="multi") try: reset_data = env.reset_multi(task_id) num_agents = reset_data["num_agents"] zone_info = reset_data["zone_info"] observations = reset_data["observations"] # Per-agent history agent_histories = {str(i): [] for i in range(num_agents)} last_rewards = {str(i): 0.0 for i in range(num_agents)} print(f"[INFO] Multi-agent mode: {num_agents} agents", flush=True) for aid, zi in zone_info.items(): print(f" Agent {aid}: {zi['zone_name']} ({len(zi['bus_ids'])} buses)", flush=True) for step_num in range(1, MAX_STEPS + 1): agent_actions = {} # Each agent generates its own action based on partial observation for agent_id_str in sorted(observations.keys()): obs = observations.get(agent_id_str, {}) zone_name = zone_info.get(agent_id_str, {}).get("zone_name", f"Zone_{agent_id_str}") obs_json = json.dumps(obs, indent=2) message = get_model_message( client, step_num, obs_json, last_rewards[agent_id_str], agent_histories[agent_id_str], SYSTEM_PROMPT_MULTI, zone_name=zone_name ) action_dict = parse_action(message) agent_actions[agent_id_str] = action_dict # Submit all actions together result = env.step_multi(agent_actions) observations = result["observations"] team_reward = result.get("team_reward", 0.0) done = result.get("done", False) # Track safety interventions safety_reports = result.get("safety_reports", {}) if isinstance(safety_reports, list): # Handle list format from older API step_interventions = sum(1 for sr in safety_reports if sr.get("was_corrected", False)) else: step_interventions = sum( 1 for sr in safety_reports.values() if sr.get("was_corrected", False) ) total_safety_interventions += step_interventions # Feed safety correction feedback into agent histories if isinstance(safety_reports, dict): for aid_str, sr in safety_reports.items(): if sr.get("was_corrected") and aid_str in agent_histories: reason = sr.get("correction_reason", "action corrected")[:120] agent_histories[aid_str].append(f"[SAFETY] {reason}") # Log per-agent rewards per_agent_rewards = result.get("rewards", {}) for agent_id_str in sorted(observations.keys()): agent_reward = per_agent_rewards.get(agent_id_str, {}).get("value", 0.0) last_rewards[agent_id_str] = agent_reward action_summary = json.dumps(agent_actions.get(agent_id_str, {})) if len(action_summary) > 100: action_summary = action_summary[:100] + "..." agent_histories[agent_id_str].append( f"Step {step_num}: action={action_summary[:60]} -> reward {agent_reward:+.2f}" ) rewards.append(team_reward) steps_taken = step_num # Log team-level step oversight = result.get("oversight_report", {}) coord_score = oversight.get("coordination_score", 1.0) safety_str = f" safety_corrections={step_interventions}" if step_interventions > 0 else "" log_step(step=step_num, action=f"team_reward={team_reward:.2f} coord={coord_score:.2f}{safety_str}", reward=team_reward, done=done) if done: break grade_result = env.grade() score = clamp_score(grade_result.get("score", 0.5)) success = score >= SUCCESS_SCORE_THRESHOLD except Exception as e: print(f"[DEBUG] Task {task_id} multi-agent error: {e}", flush=True) score = 0.05 success = False print(f"[INFO] Total safety interventions: {total_safety_interventions}", flush=True) log_end(success=success, steps=steps_taken, score=score, rewards=rewards, mode="multi") return { "task": task_id, "score": score, "steps": steps_taken, "success": success, "safety_interventions": total_safety_interventions } # ---------- Main ---------- def main(): """Run inference on all tasks.""" parser = argparse.ArgumentParser(description="OpenGrid LLM Inference") parser.add_argument("--multi", action="store_true", help="Use multi-agent POMDP mode (default: single-agent)") parser.add_argument("--tasks", nargs="+", default=TASKS, help="Which tasks to run (default: all)") args = parser.parse_args() if not API_KEY: print("[ERROR] No API key found. Set OPENAI_API_KEY or HF_TOKEN environment variable.", flush=True) sys.exit(1) mode = "multi-agent" if args.multi else "single-agent" print(f"[CONFIG] API_BASE_URL={API_BASE_URL}", flush=True) print(f"[CONFIG] MODEL_NAME={MODEL_NAME}", flush=True) print(f"[CONFIG] ENV_URL={ENV_URL}", flush=True) print(f"[CONFIG] MODE={mode}", flush=True) client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY, timeout=30.0, max_retries=2) env = EnvClient(ENV_URL) all_results = [] runner = run_task_multi if args.multi else run_task_single try: for task_id in args.tasks: print(f"\n{'='*60}", flush=True) print(f"Running task: {task_id} ({mode})", flush=True) print(f"{'='*60}", flush=True) result = runner(client, env, task_id) all_results.append(result) finally: env.close() # Summary print(f"\n{'='*60}", flush=True) print(f"FINAL RESULTS ({mode})", flush=True) print(f"{'='*60}", flush=True) for r in all_results: status = "PASS" if r["success"] else "FAIL" extra = "" if "safety_interventions" in r: extra = f" safety={r['safety_interventions']}" print(f" {r['task']}: score={r['score']:.4f} steps={r['steps']} [{status}]{extra}", flush=True) avg_score = sum(r["score"] for r in all_results) / len(all_results) if all_results else 0 print(f"\n Average Score: {avg_score:.4f}", flush=True) if __name__ == "__main__": main()