service-update-alerts/update_monitor/classifier.py

import json
import os
from pathlib import Path
from typing import Dict, List, Optional
from urllib.request import Request, urlopen


def load_update_examples(examples_dir: Path) -> List[Dict[str, str]]:
    examples = []
    for path in sorted(examples_dir.glob("*.txt")):
        text = path.read_text(encoding="utf-8")
        if "---" not in text:
            continue
        header, _, body = text.partition("---\n")
        action = reason = None
        for line in header.strip().splitlines():
            if line.startswith("action:"):
                action = line.split(":", 1)[1].strip()
            elif line.startswith("reason:"):
                reason = line.split(":", 1)[1].strip()
        if action not in ("auto", "manual"):
            continue
        examples.append({"action": action, "reason": reason or "", "notes": body.strip()})
    return examples


class Classifier:
    def __init__(
        self,
        base_url: str = "",
        model: str = "",
    ) -> None:
        self.base_url = base_url or os.getenv("OLLAMA_BASE_URL", "http://192.168.1.20:11434")
        self.model = model or os.getenv("OLLAMA_MODEL", "gemma4:26b-64k")

    def classify(
        self,
        service_dir: Path,
        release_notes: Optional[str],
        current: str,
        latest: str,
        timeout: float,
    ) -> Dict[str, str]:
        if not release_notes:
            return {"action": "manual", "reason": "No release notes available; defaulting to manual review."}

        examples_dir = service_dir / "review-examples"
        examples = load_update_examples(examples_dir) if examples_dir.is_dir() else []
        if not examples:
            return {"action": "manual", "reason": "No review examples configured; defaulting to manual review."}

        messages = self._build_messages(service_dir.name, examples, release_notes, current, latest)

        try:
            payload = json.dumps({
                "model": self.model,
                "messages": messages,
                "stream": False,
                "options": {"temperature": 0},
            }).encode("utf-8")
            req = Request(
                f"{self.base_url}/v1/chat/completions",
                data=payload,
                headers={"Content-Type": "application/json"},
                method="POST",
            )
            with urlopen(req, timeout=timeout) as resp:
                data = json.loads(resp.read())
            text = data["choices"][0]["message"]["content"].strip()
            start, end = text.find("{"), text.rfind("}") + 1
            if start == -1 or end == 0:
                raise ValueError(f"No JSON in response: {text[:200]}")
            result = json.loads(text[start:end])
            action = result.get("action", "")
            if action not in ("auto", "manual"):
                raise ValueError(f"Unexpected action value: {action!r}")
            return {"action": action, "reason": result.get("reason", "")}
        except Exception as exc:
            return {"action": "manual", "reason": f"Classification error: {exc}"}

    def _build_messages(
        self,
        service_name: str,
        examples: List[Dict[str, str]],
        release_notes: str,
        current: str,
        latest: str,
    ) -> List[Dict[str, str]]:
        system_prompt = (
            f"You are an update safety classifier for self-hosted {service_name} instances.\n"
            "Given release notes for a software update, decide whether the update can be applied\n"
            "automatically (action: auto) or requires the admin to intervene manually before or\n"
            "during the update (action: manual).\n\n"
            "Rules:\n"
            "- Classify only THIS release, not future or past releases mentioned in the notes.\n"
            "- Warnings or announcements about breaking changes in upcoming versions do NOT\n"
            "  make the current release manual.\n"
            "- Conditional notes like 'follow X instructions if upgrading from <= Y' only apply\n"
            "  if the current running version is <= Y. If we are already past Y, they do not apply.\n"
            '- Respond ONLY with a JSON object: {"action": "auto", "reason": "..."}\n'
            '  or {"action": "manual", "reason": "..."}\n'
            "- Keep reason under 60 words."
        )
        messages: List[Dict[str, str]] = [{"role": "system", "content": system_prompt}]
        for ex in examples:
            messages.append({"role": "user", "content": f"Release notes:\n{ex['notes']}"})
            messages.append({
                "role": "assistant",
                "content": json.dumps({"action": ex["action"], "reason": ex["reason"]}),
            })
        messages.append({
            "role": "user",
            "content": f"Upgrading from {current} to {latest}.\n\nRelease notes:\n{release_notes}",
        })
        return messages