client.py

#!/usr/bin/env python3
"""FixOS baseline inference script."""

import json
import os
import sys
import time
from typing import Any, Dict, Optional
from urllib import error as urlerror
from urllib import request as urlrequest

from openai import APIConnectionError, APIError, OpenAI, RateLimitError


def _load_local_env_class():
    errors = []
    try:
        from server.my_env_environment import FixOSEnvironment  # type: ignore

        return FixOSEnvironment, ""
    except Exception as exc:
        errors.append(f"server.my_env_environment: {exc}")

    try:
        from my_env.server.my_env_environment import FixOSEnvironment  # type: ignore

        return FixOSEnvironment, ""
    except Exception as exc:
        errors.append(f"my_env.server.my_env_environment: {exc}")

    return None, " | ".join(errors)

API_BASE_URL = os.getenv("API_BASE_URL", "https://api.openai.com/v1")
MODEL_NAME = os.getenv("MODEL_NAME", "gpt-4o-mini")


def _env_required(name: str) -> str:
    value = os.getenv(name)
    if not value:
        raise ValueError(f"Missing required environment variable: {name}")
    return value


class EnvHTTPClient:
    def __init__(self, base_url: str):
        self.base_url = base_url.rstrip("/")

    def reset(self) -> Dict[str, Any]:
        return self._post_json("/reset", {})

    def step(self, action: Dict[str, Any]) -> Dict[str, Any]:
        try:
            return self._post_json("/step", action)
        except RuntimeError as exc:
            if "422" in str(exc):
                return self._post_json("/step", {"action": action})
            raise

    def _post_json(self, path: str, payload: Dict[str, Any]) -> Dict[str, Any]:
        req = urlrequest.Request(
            url=f"{self.base_url}{path}",
            data=json.dumps(payload).encode("utf-8"),
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        try:
            with urlrequest.urlopen(req, timeout=30) as resp:
                return json.loads(resp.read().decode("utf-8"))
        except urlerror.URLError as exc:
            raise RuntimeError(f"HTTP request failed for {path}: {exc}") from exc


class LocalEnvClient:
    def __init__(self):
        FixOSEnvironment, import_error = _load_local_env_class()
        if FixOSEnvironment is None:
            raise RuntimeError(f"Local FixOS environment is unavailable (import failed): {import_error}")
        self.env = FixOSEnvironment()

    def reset(self) -> Dict[str, Any]:
        observation = self.env.reset()
        return {"observation": observation.model_dump(), "reward": 0.0, "done": False}

    def step(self, action: Dict[str, Any]) -> Dict[str, Any]:
        try:
            from models import FixOSAction
        except ImportError:
            from my_env.models import FixOSAction

        observation = self.env.step(FixOSAction(command=action.get("command", "status"), args=action.get("args", {})))
        return {
            "observation": observation.model_dump(),
            "reward": observation.reward,
            "done": observation.done,
        }


class FixOSAgent:
    def __init__(self, env_url: str | None = None):
        hf_token = _env_required("HF_TOKEN")

        self.llm = OpenAI(api_key=hf_token, base_url=API_BASE_URL)
        self.model_name = MODEL_NAME
        self.env = self._create_env_client(env_url)

    def _create_env_client(self, env_url: str | None):
        if not env_url:
            return LocalEnvClient()
        try:
            client = EnvHTTPClient(env_url)
            client.reset()
            return client
        except Exception:
            return LocalEnvClient()

    def _emit(self, tag: str, payload: Dict[str, Any]) -> None:
        print(f"[{tag}] {json.dumps(payload, separators=(',', ':'), sort_keys=False)}", flush=True)

    def _llm_action(self, observation: Dict[str, Any], step: int, max_steps: int) -> Dict[str, Any]:
        prompt = (
            "You are solving a deterministic OS troubleshooting task.\n"
            "Return JSON only with keys: command, args, reasoning.\n"
            f"Step: {step}/{max_steps}\n"
            f"Observation: {json.dumps(observation)}\n"
            "Allowed commands: ps, top, df, status, logs, cat, edit, restart, kill, rm.\n"
            "Prefer concrete remediation over repeated diagnostics."
        )

        try:
            resp = self.llm.chat.completions.create(
                model=self.model_name,
                messages=[{"role": "user", "content": prompt}],
                temperature=0.0,
                max_tokens=250,
            )
            content = (resp.choices[0].message.content or "").strip()
            start = content.find("{")
            end = content.rfind("}")
            if start >= 0 and end >= start:
                candidate = json.loads(content[start : end + 1])
                cmd = str(candidate.get("command", "status")).lower()
                args = candidate.get("args", {})
                if cmd in {"ps", "top", "df", "status", "logs", "cat", "edit", "restart", "kill", "rm"} and isinstance(args, dict):
                    return {"command": cmd, "args": args, "reasoning": str(candidate.get("reasoning", ""))}
        except (APIError, APIConnectionError, RateLimitError, ValueError):
            pass

        return self._heuristic_action(observation)

    def _heuristic_action(self, observation: Dict[str, Any]) -> Dict[str, Any]:
        services = {s.get("name", ""): s for s in observation.get("services", [])}
        resources = observation.get("resources", {})
        processes = observation.get("processes", [])
        history = observation.get("history", [])

        high_cpu = sorted(processes, key=lambda p: p.get("cpu_percent", 0), reverse=True)

        for proc in high_cpu:
            if int(proc.get("pid", -1)) == 922:
                return {"command": "kill", "args": {"pid": 922}, "reasoning": "remove port blocker"}

        if float(resources.get("cpu_percent", 0)) > 80:
            for proc in high_cpu:
                if int(proc.get("pid", -1)) in {920, 921}:
                    return {
                        "command": "kill",
                        "args": {"pid": int(proc.get("pid"))},
                        "reasoning": "reduce aggregate cpu pressure",
                    }

        if float(resources.get("disk_percent", 0)) > 95:
            for candidate in ["/var/log/archive.bin", "/var/log/system.log"]:
                if any(f.get("path") == candidate for f in observation.get("filesystem", [])):
                    return {"command": "rm", "args": {"path": candidate}, "reasoning": "free disk"}

        for proc in high_cpu:
            if float(proc.get("cpu_percent", 0)) >= 50:
                return {"command": "kill", "args": {"pid": int(proc.get("pid"))}, "reasoning": "kill high cpu process"}

        nginx = services.get("nginx", {})
        mysql = services.get("mysql", {})

        if nginx and not nginx.get("config_valid", True):
            return {"command": "edit", "args": {"path": "/etc/nginx/nginx.conf", "content": "valid nginx config"}, "reasoning": "fix nginx config"}

        if mysql and not mysql.get("config_valid", True):
            return {"command": "edit", "args": {"path": "/etc/mysql/my.cnf", "content": "valid mysql config"}, "reasoning": "fix mysql config"}

        if mysql.get("status") != "running":
            return {"command": "restart", "args": {"service": "mysql"}, "reasoning": "restart mysql"}

        if nginx.get("status") != "running":
            return {"command": "restart", "args": {"service": "nginx"}, "reasoning": "restart nginx"}

        recent = " ".join(str(item) for item in history[-3:]).lower()
        if "logs" not in recent and "cat" not in recent:
            return {"command": "logs", "args": {}, "reasoning": "check logs"}

        return {"command": "status", "args": {}, "reasoning": "verify system"}

    def run_episode(
        self,
        episode_index: int,
        max_steps: int = 50,
        retried_local: bool = False,
        emit_start: bool = True,
    ) -> Dict[str, Any]:
        try:
            reset_payload = self.env.reset()
        except Exception:
            if not retried_local:
                self.env = LocalEnvClient()
                return self.run_episode(
                    episode_index,
                    max_steps=max_steps,
                    retried_local=True,
                    emit_start=emit_start,
                )
            raise

        obs = reset_payload.get("observation", {})
        task_id = str(obs.get("task_id", "unknown"))
        episode_id = f"ep-{episode_index:03d}"

        if emit_start:
            self._emit(
                "START",
                {
                    "episode_id": episode_id,
                    "task_id": task_id,
                    "max_steps": max_steps,
                    "timestamp": int(time.time()),
                },
            )

        total_reward = 0.0
        success = False

        for step in range(1, max_steps + 1):
            action_obj = self._llm_action(obs, step, max_steps)
            try:
                step_payload = self.env.step({"command": action_obj["command"], "args": action_obj["args"]})
            except Exception:
                if not retried_local:
                    self.env = LocalEnvClient()
                    return self.run_episode(
                        episode_index,
                        max_steps=max_steps,
                        retried_local=True,
                        emit_start=False,
                    )
                raise
            obs = step_payload.get("observation", {})

            reward = float(step_payload.get("reward", obs.get("reward", 0.0) or 0.0))
            done = bool(step_payload.get("done", obs.get("done", False)))
            total_reward += reward
            success = bool(obs.get("is_success_step", False)) or success

            self._emit(
                "STEP",
                {
                    "episode_id": episode_id,
                    "task_id": task_id,
                    "step": step,
                    "command": action_obj["command"],
                    "args": action_obj["args"],
                    "reward": round(reward, 6),
                    "task_score": round(float(obs.get("task_score", 0.0) or 0.0), 6),
                    "done": done,
                },
            )

            if done:
                break

        final_score = max(0.0, min(1.0, float(obs.get("task_score", 0.0) or 0.0)))
        self._emit(
            "END",
            {
                "episode_id": episode_id,
                "task_id": task_id,
                "success": success,
                "steps_taken": len(obs.get("history", [])),
                "total_reward": round(total_reward, 6),
                "final_score": round(final_score, 6),
                "timestamp": int(time.time()),
            },
        )

        return {"task_id": task_id, "success": success, "score": final_score}


def main() -> None:
    agent = FixOSAgent(env_url=os.getenv("ENV_BASE_URL"))

    runs = 7
    scores: Dict[str, list[float]] = {}

    for i in range(runs):
        result = agent.run_episode(i + 1, max_steps=50)
        scores.setdefault(result["task_id"], []).append(float(result["score"]))

    summary = {k: round(sum(v) / len(v), 6) for k, v in sorted(scores.items())}
    print(json.dumps({"summary": summary}, separators=(",", ":")), file=sys.stderr)


if __name__ == "__main__":
    main()