agent/baseline_agent.py

"""
Baseline Agent — Rule-based agent for validating the RL loop end-to-end.

This agent uses hardcoded heuristics per error type to generate fix commands.
Use it to confirm the environment, executor, and reward engine are working
correctly BEFORE plugging in the LLM. Per the hackathon guide: "Start from
a capable model and validate the loop before RL."
"""

from __future__ import annotations

import re
from typing import Dict, List


class BaselineAgent:
    """Deterministic rule-based DevOps troubleshooting agent.

    Maps error types to predefined fix strategies using pattern matching.
    Serves as:
      1. Loop validator — confirms reset/step/reward work end-to-end
      2. Performance baseline — the LLM agent must beat this
      3. Fallback — used when no GPU/LLM is available

    Usage:
        agent = BaselineAgent()
        command = agent.act(observation)
    """

    def act(self, observation: Dict) -> str:
        """Generate a shell command from the current observation.

        Routes to a strategy handler based on the error_type field
        in the observation, then uses command_history to avoid repeats.

        Args:
            observation: Dict with error_log, command_history, error_type,
                         step_count, scenario_id.

        Returns:
            Shell command string.
        """
        error_log = observation.get("error_log", "")
        error_type = observation.get("error_type", "unknown")
        history = observation.get("command_history", [])
        scenario_id = observation.get("scenario_id", "")

        # First try scenario-specific strategies (most precise)
        scenario_cmd = self._scenario_strategy(scenario_id, error_log, history)
        if scenario_cmd:
            return scenario_cmd

        # Then fall back to error-type strategies
        handlers = {
            "missing_package": self._fix_missing_package,
            "port_conflict": self._fix_port_conflict,
            "missing_env": self._fix_missing_env,
            "version_conflict": self._fix_version_conflict,
            "syntax_error": self._fix_syntax_error,
            "config_error": self._fix_config_error,
            "file_not_found": self._fix_file_not_found,
            "service_not_running": self._fix_service_not_running,
            "permission_denied": self._fix_permission_denied,
        }

        handler = handlers.get(error_type, self._fix_unknown)
        return handler(error_log, history)

    def _scenario_strategy(self, scenario_id: str, error_log: str, history: List[str]) -> str | None:
        """Try scenario-specific optimal strategies first.

        Args:
            scenario_id: The scenario identifier.
            error_log: Current error log.
            history: Previous commands this episode.

        Returns:
            Command string or None if no specific strategy.
        """
        strategies = {
            "missing_flask": ["pip install flask", "python /app/main.py"],
            "missing_numpy": ["pip install numpy", "python /app/main.py"],
            "missing_requests": ["pip install requests", "python /app/main.py"],
            "wrong_python_version": ["python3 /app/main.py"],
            "port_conflict": [
                "lsof -t -i:5000 | xargs kill -9",
                "python /app/server.py &",
            ],
            "missing_env_var": [
                "export DATABASE_URL=postgresql://localhost:5432/mydb",
                "python /app/db_app.py",
            ],
            "broken_requirements": [
                "sed -i 's/werkzeug==1.0.0/werkzeug>=2.3.0/' /app/requirements.txt",
                "pip install -r /app/requirements.txt",
                "python /app/main.py",
            ],
            "corrupt_venv": [
                "rm -rf /app/venv",
                "python3 -m venv /app/venv",
                "source /app/venv/bin/activate && pip install flask",
                "source /app/venv/bin/activate && python -c 'import flask; print(\"Success\")'",
            ],
            "wrong_config_restart": [
                "sed -i 's/127.0.0.1/0.0.0.0/' /app/config.py",
                "kill $(lsof -t -i:8080) 2>/dev/null; true",
                "python /app/server.py &",
            ],
            "db_migration_fail": [
                "cat > /app/models.py << 'PYEOF'\nfrom sqlalchemy import Column, Integer, String, create_engine\nfrom sqlalchemy.orm import declarative_base\nBase = declarative_base()\nclass User(Base):\n    __tablename__ = 'users'\n    id = Column(Integer, primary_key=True)\n    name = Column(String(100), nullable=False)\n    email = Column(String(200), unique=True)\nPYEOF",
                "pip install sqlalchemy -q",
                "python /app/migrate.py",
            ],
        }

        if scenario_id not in strategies:
            return None

        cmds = strategies[scenario_id]
        step = len(history)
        if step < len(cmds):
            return cmds[step]

        # All hint commands exhausted — try verification
        return None

    def _fix_missing_package(self, error_log: str, history: List[str]) -> str:
        """Handle ModuleNotFoundError / ImportError."""
        match = re.search(r"No module named ['\"]?(\w+)", error_log)
        if match:
            module = match.group(1)
            cmd = f"pip install {module}"
            if cmd not in history:
                return cmd
            return f"pip3 install {module}"
        return "pip install -r requirements.txt"

    def _fix_port_conflict(self, error_log: str, history: List[str]) -> str:
        """Handle Address already in use."""
        match = re.search(r"port\s+(\d+)", error_log, re.IGNORECASE)
        if not match:
            match = re.search(r":(\d{4,5})", error_log)
        port = match.group(1) if match else "5000"

        if not any("kill" in cmd for cmd in history):
            return f"lsof -t -i:{port} | xargs kill -9"
        return "python /app/server.py &"

    def _fix_missing_env(self, error_log: str, history: List[str]) -> str:
        """Handle KeyError on environment variables."""
        match = re.search(r"KeyError:\s*['\"](\w+)['\"]", error_log)
        if match:
            var = match.group(1)
            defaults = {
                "DATABASE_URL": "postgresql://localhost:5432/mydb",
                "SECRET_KEY": "dev-secret-key-12345",
                "API_KEY": "test-api-key",
                "REDIS_URL": "redis://localhost:6379",
            }
            if not any("export" in cmd for cmd in history):
                value = defaults.get(var, "placeholder_value")
                return f"export {var}={value}"
            # Env var set; now rerun the app
            match_file = re.search(r'File "([^"]+)"', error_log)
            if match_file:
                return f"python {match_file.group(1)}"
            return "python /app/db_app.py"
        return "env"

    def _fix_version_conflict(self, error_log: str, history: List[str]) -> str:
        """Handle package version conflicts."""
        if not any("sed" in cmd for cmd in history):
            match = re.search(r"requested\s+(\w+)==(\S+)", error_log)
            if match:
                pkg = match.group(1)
                return f"sed -i 's/{pkg}==.*/{pkg}>=0/' /app/requirements.txt"
            return "sed -i 's/werkzeug==1.0.0/werkzeug>=2.3.0/' /app/requirements.txt"
        return "pip install -r /app/requirements.txt"

    def _fix_syntax_error(self, error_log: str, history: List[str]) -> str:
        """Handle SyntaxError (usually python version mismatch)."""
        if "python2" in error_log.lower() or "shebang" in error_log.lower():
            match = re.search(r'File "([^"]+)"', error_log)
            if match:
                return f"python3 {match.group(1)}"
        return "python3 /app/main.py"

    def _fix_config_error(self, error_log: str, history: List[str]) -> str:
        """Handle configuration errors."""
        if "127.0.0.1" in error_log:
            if not any("sed" in cmd for cmd in history):
                return "sed -i 's/127.0.0.1/0.0.0.0/' /app/config.py"
            if not any("kill" in cmd for cmd in history):
                return "kill $(lsof -t -i:8080) 2>/dev/null; true"
            return "python /app/server.py &"

        if "NameError" in error_log or "INVALID" in error_log:
            match = re.search(r'File "([^"]+)"', error_log)
            if match and not any("cat >" in cmd for cmd in history):
                return f"cat {match.group(1)}"
            return "python /app/migrate.py"

        return "cat /app/config.py"

    def _fix_file_not_found(self, error_log: str, history: List[str]) -> str:
        """Handle FileNotFoundError / No such file or directory."""
        if "venv" in error_log or "bad interpreter" in error_log:
            if not any("rm" in cmd for cmd in history):
                return "rm -rf /app/venv"
            if not any("python3 -m venv" in cmd for cmd in history):
                return "python3 -m venv /app/venv"
            return "source /app/venv/bin/activate && pip install flask"

        match = re.search(r"No such file.*?['\"]?(/\S+?)(?:['\"]|$)", error_log)
        if match:
            return f"ls -la {match.group(1)}"
        return "ls -la /app/"

    def _fix_service_not_running(self, error_log: str, history: List[str]) -> str:
        """Handle Connection refused / service not running."""
        match = re.search(r"port\s+(\d+)", error_log, re.IGNORECASE)
        if not match:
            match = re.search(r":(\d{4,5})", error_log)
        port = match.group(1) if match else "8080"
        return f"python /app/server.py --port {port} &"

    def _fix_permission_denied(self, error_log: str, history: List[str]) -> str:
        """Handle PermissionError."""
        match = re.search(r"'([^']+)'", error_log)
        if match:
            return f"chmod +x {match.group(1)}"
        return "ls -la /app/"

    def _fix_unknown(self, error_log: str, history: List[str]) -> str:
        """Fallback for unclassified errors."""
        if not history:
            return "ls -la /app/"
        if len(history) == 1:
            return "cat /app/*.py 2>/dev/null || echo 'no python files'"
        return "echo 'Unable to determine fix strategy'"