Spaces:
Running
Running
| import json | |
| import os | |
| import re | |
| from typing import Dict, Any, Optional, List | |
| from openai import OpenAI | |
| from .environment import TrainingEnv | |
| class AIWorker: | |
| """Programmatic agent interface — drive the environment step by step.""" | |
| def __init__(self, scenario: str = "log_analysis"): | |
| self.engine = TrainingEnv(scenario=scenario) | |
| self.recorder: List[Dict[str, Any]] = [] | |
| def boot(self) -> Dict[str, Any]: | |
| reset_out = self.engine.reset() | |
| return { | |
| "status": "reset", | |
| "observation": reset_out["observation"], | |
| "task_name": reset_out["info"]["task_name"], | |
| "task_instructions": reset_out["info"]["instructions"], | |
| "max_steps": reset_out["info"]["max_steps"], | |
| "message": "Environment reset. Read the task instructions and begin solving." | |
| } | |
| def invoke(self, action: str, rationale: Optional[str] = None) -> Dict[str, Any]: | |
| step_out = self.engine.step(action) | |
| resp = { | |
| "status": "success" if step_out["info"]["exit_code"] == 0 else "failed", | |
| "command": action, | |
| "command_output": step_out["info"].get("command_output", ""), | |
| "reward": step_out["reward"], | |
| "task_score": step_out["info"]["task_score"], | |
| "done": step_out["done"], | |
| "observation": step_out["observation"], | |
| "step": step_out["info"]["step"], | |
| "max_steps": step_out["info"]["max_steps"], | |
| "task_metadata": step_out["info"]["task_metadata"], | |
| } | |
| self.recorder.append({ | |
| "command": action, | |
| "reasoning": rationale, | |
| "result": resp, | |
| }) | |
| if resp["done"]: | |
| resp["message"] = f"Task completed! Final score: {resp['task_score']:.1f}" | |
| else: | |
| resp["message"] = f"Score: {resp['task_score']:.1f} | Steps: {resp['step']}/{resp['max_steps']}" | |
| return resp | |
| def context(self) -> str: | |
| state = self.engine.dump() | |
| recent = self.recorder[-5:] if self.recorder else [] | |
| recent_cmds = "\n".join( | |
| f" [{r['command']}] -> score={r['result']['task_score']}" | |
| for r in recent | |
| ) or "No commands executed yet" | |
| return ( | |
| f"LINUX SRE ENVIRONMENT STATE:\n" | |
| f"=== CURRENT TASK ===\n" | |
| f"Name: {self.engine.task.nm}\n" | |
| f"Difficulty: {self.engine.difficulty}\n" | |
| f"Description: {self.engine.task.desc}\n\n" | |
| f"=== PROGRESS ===\n" | |
| f"Task Score: {self.engine.score:.1f}\n" | |
| f"Steps Taken: {self.engine.step_count}\n" | |
| f"Max Steps: {self.engine.limit}\n\n" | |
| f"=== RECENT COMMANDS ===\n{recent_cmds}\n\n" | |
| f"=== WORKING DIRECTORY ===\n{self.engine.terminal.cwd}\n" | |
| ) | |
| def report(self) -> Dict[str, Any]: | |
| return { | |
| "task_name": self.engine.task.nm, | |
| "difficulty": self.engine.difficulty, | |
| "final_score": self.engine.score, | |
| "steps_used": self.engine.step_count, | |
| "max_steps": self.engine.limit, | |
| "completed": self.engine.finished, | |
| "efficiency": round(1 - (self.engine.step_count / self.engine.limit), 2) | |
| if self.engine.finished else 0.0, | |
| "command_history": self.engine.terminal.history(), | |
| } | |
| # ====================================================================== | |
| # LLM AGENT (OpenAI client) | |
| # ====================================================================== | |
| class LLMAgent: | |
| """ | |
| LLM-powered agent that autonomously solves scenarios. | |
| Uses the OpenAI client directly and reads configuration from: | |
| - API_BASE_URL | |
| - MODEL_NAME | |
| - HF_TOKEN | |
| Usage: | |
| agent = LLMAgent() | |
| result = agent.solve("cascading_db_failure") | |
| """ | |
| def __init__( | |
| self, | |
| model: Optional[str] = None, | |
| api_key: Optional[str] = None, | |
| base_url: Optional[str] = None, | |
| max_turns: int = 30, | |
| verbose: bool = False, | |
| ): | |
| self.model = self._normalize_model_name(model or os.environ.get("MODEL_NAME", "")) | |
| self.api_key = api_key or os.environ.get("HF_TOKEN", "") | |
| self.base_url = base_url or os.environ.get("API_BASE_URL", "") | |
| self.max_turns = max_turns | |
| self.verbose = verbose | |
| self._client: Optional[OpenAI] = None | |
| def _normalize_model_name(model_name: str) -> str: | |
| if "/" in model_name: | |
| provider, candidate = model_name.split("/", 1) | |
| if provider in {"openai", "anthropic", "gemini"}: | |
| return candidate | |
| return model_name | |
| def _ensure_client(self) -> OpenAI: | |
| if self._client is None: | |
| missing = [] | |
| if not self.base_url: | |
| missing.append("API_BASE_URL") | |
| if not self.model: | |
| missing.append("MODEL_NAME") | |
| if not self.api_key: | |
| missing.append("HF_TOKEN") | |
| if missing: | |
| raise RuntimeError( | |
| "Missing required environment variables for LLM agent: " | |
| + ", ".join(missing) | |
| ) | |
| kwargs: Dict[str, Any] = {"api_key": self.api_key, "base_url": self.base_url} | |
| self._client = OpenAI(**kwargs) | |
| return self._client | |
| def _build_messages(self, initial: Dict[str, Any]) -> List[Dict[str, str]]: | |
| return [ | |
| { | |
| "role": "system", | |
| "content": ( | |
| SystemPrompts.get_sys(initial.get("task_name", "linux-sre")) | |
| + "\n\nRespond with a short explanation and a single command." | |
| + "\nPreferred format: ```bash\\n<command>\\n```" | |
| ), | |
| }, | |
| { | |
| "role": "user", | |
| "content": SystemPrompts.format_observation(initial), | |
| }, | |
| ] | |
| def solve(self, scenario: str = "log_analysis") -> Dict[str, Any]: | |
| """Run the agent loop: observe, think, act, and repeat until done.""" | |
| client = self._ensure_client() | |
| worker = AIWorker(scenario=scenario) | |
| initial = worker.boot() | |
| messages = self._build_messages(initial) | |
| turns = [] | |
| for turn in range(self.max_turns): | |
| response = client.chat.completions.create( | |
| model=self.model, | |
| messages=messages, | |
| temperature=0.2, | |
| max_tokens=512, | |
| ) | |
| assistant_msg = (response.choices[0].message.content or "").strip() | |
| messages.append({"role": "assistant", "content": assistant_msg}) | |
| command = self._extract_command(assistant_msg) | |
| if not command: | |
| messages.append({"role": "user", "content": "Return exactly one shell command."}) | |
| continue | |
| result = worker.invoke(command, rationale=assistant_msg) | |
| turns.append({"turn": turn + 1, "command": command, "score": result["task_score"]}) | |
| if self.verbose: | |
| print(f"[Turn {turn + 1}] $ {command} -> score={result['task_score']:.1f}") | |
| if result["done"]: | |
| break | |
| messages.append( | |
| { | |
| "role": "user", | |
| "content": ( | |
| f"Command output:\n{result['command_output']}\n\n" | |
| f"Score: {result['task_score']:.1f} | Step {result['step']}/{result['max_steps']}\n\n" | |
| "What command should we run next?" | |
| ), | |
| } | |
| ) | |
| summary = worker.report() | |
| summary["turns"] = turns | |
| summary["model"] = self.model | |
| summary["base_url"] = self.base_url | |
| return summary | |
| def extract_command(text: str) -> Optional[str]: | |
| """Extract a shell command from an LLM response.""" | |
| try: | |
| payload = json.loads(text) | |
| if isinstance(payload, dict): | |
| command = payload.get("command") | |
| if isinstance(command, str) and command.strip(): | |
| return command.strip() | |
| except Exception: | |
| pass | |
| fenced = re.search(r"```(?:bash|sh|shell)?\s*\n(.+?)\n```", text, re.DOTALL) | |
| if fenced: | |
| for line in fenced.group(1).strip().split("\n"): | |
| line = line.strip() | |
| if line and not line.startswith("#"): | |
| return line.lstrip("$ ") | |
| inline = re.search(r"`([^`]+)`", text) | |
| if inline: | |
| command = inline.group(1).strip() | |
| if command: | |
| return command | |
| for line in text.split("\n"): | |
| line = line.strip() | |
| if line.startswith("$ "): | |
| return line[2:] | |
| return None | |
| def _extract_command(text: str) -> Optional[str]: | |
| return LLMAgent.extract_command(text) | |
| # ====================================================================== | |
| # PROMPT TEMPLATES | |
| # ====================================================================== | |
| class SystemPrompts: | |
| def get_sys(task_name: str) -> str: | |
| return ( | |
| "You are a Linux System Reliability Engineer (SRE) agent. " | |
| "You are troubleshooting a production incident.\n\n" | |
| "RULES:\n" | |
| "- Execute ONE command at a time\n" | |
| "- Wrap your command in a code block: ```bash\\ncommand\\n```\n" | |
| "- Briefly explain your reasoning before each command\n" | |
| "- Work systematically: diagnose first, then fix\n" | |
| "- Verify your fixes after applying them\n\n" | |
| "Available commands: cat, grep, ls, chmod, ps, kill, systemctl, " | |
| "cd, tail, head, find, wc, df, du, top, free, uptime, whoami, " | |
| "id, hostname, date, touch, mkdir, cp, mv, rm, netstat, curl, " | |
| "journalctl, crontab, iptables, mount, env, export, sort, uniq, " | |
| "cut, tr, tee, echo, pwd, dig\n\n" | |
| "Pipes (|) and redirects (>, >>) are supported.\n\n" | |
| f"Current task: {task_name}" | |
| ) | |
| def format_observation(initial: Dict[str, Any]) -> str: | |
| return ( | |
| f"TASK: {initial['task_name']}\n\n" | |
| f"INSTRUCTIONS:\n{initial['task_instructions']}\n\n" | |
| f"CURRENT STATE:\n" | |
| f"- Directory: {initial['observation']['current_directory']}\n" | |
| f"- Processes:\n{initial['observation']['processes']}\n\n" | |
| f"You have {initial['max_steps']} steps. Begin troubleshooting." | |
| ) | |
| # ====================================================================== | |
| # DEMO | |
| # ====================================================================== | |
| def demo(): | |
| """Quick demo showing the agent solving the medium task.""" | |
| agent = AIWorker(scenario="permission_repair") | |
| initial = agent.boot() | |
| print(f"Task: {initial['task_name']}") | |
| print(f"Scripts: {initial['observation']['filesystem']}") | |
| commands = [ | |
| ("ls -la /home/user/scripts/cleanup.sh", "Check current permissions"), | |
| ("chmod 0755 /home/user/scripts/cleanup.sh", "Make script executable"), | |
| ("ls -la /home/user/scripts/cleanup.sh", "Verify permissions fixed"), | |
| ] | |
| for cmd, reason in commands: | |
| result = agent.invoke(cmd, rationale=reason) | |
| print(f"\n$ {cmd} ({reason})") | |
| print(f" Output: {result['command_output'][:120]}") | |
| print(f" Score: {result['task_score']:.1f} | Done: {result['done']}") | |
| if result["done"]: | |
| break | |
| print(f"\nFinal report: {json.dumps(agent.report(), indent=2)}") | |
| if __name__ == "__main__": | |
| demo() | |