LLM-Zabbix-HABR/alert-processor/app/llm_remediation.py

from __future__ import annotations

import json
import re
from dataclasses import dataclass
from typing import Any

import httpx

from app.models import NotificationDecision, ProcessorForwardEnvelope


@dataclass
class RemediationResult:
    ok: bool
    summary: str | None
    steps: list[str]
    commands: list[str]
    error: str | None = None


class LLMRemediationAdapter:
    def __init__(
        self,
        base_url: str,
        model: str,
        timeout_seconds: float = 45,
        verify_tls: bool = True,
        temperature: float = 0.1,
        max_steps: int = 4,
        max_commands: int = 4,
    ) -> None:
        self.base_url = base_url.rstrip("/")
        self.model = model
        self.temperature = temperature
        self.max_steps = max_steps
        self.max_commands = max_commands

        self.client = httpx.AsyncClient(
            timeout=timeout_seconds,
            verify=verify_tls,
        )

    async def close(self) -> None:
        await self.client.aclose()

    async def generate(
        self,
        envelope: ProcessorForwardEnvelope,
        decision: NotificationDecision,
    ) -> RemediationResult:
        prompt = self._build_prompt(envelope, decision)

        response = await self.client.post(
            f"{self.base_url}/api/generate",
            json={
                "model": self.model,
                "prompt": prompt,
                "stream": False,
                "format": "json",
                "think": False,
                "options": {
                    "temperature": self.temperature,
                },
            },
        )
        response.raise_for_status()

        data = response.json()
        raw_text = data.get("response", "")

        parsed = self._parse_json_object(raw_text)

        summary = self._normalize_text(parsed.get("summary"))
        steps = self._normalize_list(parsed.get("steps"), self.max_steps)
        commands = self._normalize_list(parsed.get("commands"), self.max_commands)

        # Пост-фильтрация команд
        alert_scope = (
            envelope.event.zabbix_context.get("alert_scope")
            if envelope.event.zabbix_context
            else "unknown"
        )
        commands = self._post_filter_commands(commands, alert_scope=alert_scope)

        if not summary and not steps and not commands:
            return RemediationResult(
                ok=False,
                summary=None,
                steps=[],
                commands=[],
                error="LLM returned empty remediation payload",
            )

        return RemediationResult(
            ok=True,
            summary=summary,
            steps=steps,
            commands=commands,
            error=None,
        )

    def _build_prompt(
        self,
        envelope: ProcessorForwardEnvelope,
        decision: NotificationDecision,
    ) -> str:
        event = envelope.event
        zbx_context = event.zabbix_context or {}

        context = {
            "severity": decision.severity,
            "reason": decision.reason,
            "event_phase": decision.event_phase,
            "host": event.host,
            "service": event.service,
            "trigger_name": event.trigger_name,
            "fingerprint": decision.fingerprint,
            "repeat_count": decision.repeat_count,
            "item_id": event.item_id,
            "trigger_id": event.trigger_id,
            "event_id": event.event_id,
            "value": event.value,
            "tags": event.tags,
            "zabbix_context": zbx_context,
            "alert_scope": zbx_context.get("alert_scope", "unknown"),
        }

        return (
            "You are a cautious SRE assistant. /no_think\n"
            "Generate remediation guidance for a Zabbix alert.\n"
            "Rules:\n"
            "1. Output JSON only.\n"
            "2. Do not suggest destructive commands.\n"
            "3. Commands must be diagnostic or read-only.\n"
            "4. Keep it concise.\n"
            "5. If this is a recovery event, return empty remediation.\n"
            "6. Do not assume that host name indicates container name.\n"
            "7. If alert_scope is host_os, diagnose the Linux host, not a container.\n"
            "8. Only suggest docker/container commands when alert_scope is container.\n"
            "9. Prefer generic Linux diagnostic commands for host_os alerts.\n"
            "10. Avoid restart, rm, kill -9, stop, prune, delete, drop, truncate, format.\n\n"
            "Return exactly this JSON schema:\n"
            "{\n"
            '  "summary": "short explanation",\n'
            '  "steps": ["step1", "step2"],\n'
            '  "commands": ["cmd1", "cmd2"]\n'
            "}\n\n"
            f"Alert context:\n{json.dumps(context, ensure_ascii=False, indent=2)}"
        )

    @staticmethod
    def _parse_json_object(raw_text: str) -> dict[str, Any]:
        text = raw_text.strip()

        try:
            return json.loads(text)
        except json.JSONDecodeError:
            pass

        fenced = re.search(r"\{.*\}", text, flags=re.DOTALL)
        if fenced:
            return json.loads(fenced.group(0))

        raise ValueError(f"Unable to parse LLM JSON response: {raw_text}")

    @staticmethod
    def _normalize_text(value: Any) -> str | None:
        if value is None:
            return None
        text = str(value).strip()
        return text or None

    @staticmethod
    def _normalize_list(value: Any, limit: int) -> list[str]:
        if value is None:
            return []

        if isinstance(value, list):
            result = []
            for item in value:
                text = str(item).strip()
                if text:
                    result.append(text)
            return result[:limit]

        text = str(value).strip()
        return [text] if text else []

    @staticmethod
    def _post_filter_commands(commands: list[str], alert_scope: str) -> list[str]:
        destructive_patterns = (
            r"\brm\b",
            r"\bmv\b",
            r"\bshutdown\b",
            r"\breboot\b",
            r"\bpoweroff\b",
            r"\bhalt\b",
            r"\bmkfs\b",
            r"\bfdisk\b",
            r"\bdd\b",
            r"\bkill\s+-9\b",
            r"\bdocker\s+rm\b",
            r"\bdocker\s+stop\b",
            r"\bdocker\s+restart\b",
            r"\bdocker\s+system\s+prune\b",
            r"\bkubectl\s+delete\b",
            r"\btruncate\b",
            r"\bdrop\b",
        )

        container_patterns = (
            r"\bdocker\b",
            r"\bpodman\b",
            r"\bkubectl\b",
            r"\bcrictl\b",
        )

        filtered: list[str] = []

        for cmd in commands:
            normalized = cmd.strip()
            if not normalized:
                continue

            lowered = normalized.lower()

            # Убираем деструктивные команды
            if any(re.search(pattern, lowered) for pattern in destructive_patterns):
                continue

            # Если это host_os, выбрасываем container-специфичные команды
            if alert_scope == "host_os":
                if any(re.search(pattern, lowered) for pattern in container_patterns):
                    continue

            filtered.append(normalized)

        return filtered