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 = "",
        fallback_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")
        self.fallback_model = fallback_model or os.getenv("OLLAMA_FALLBACK_MODEL", "gemma4:12b")
        self.api_key = os.getenv("OLLAMA_API_KEY", "")

    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."}

        release_notes = self._truncate(release_notes)
        context_file = service_dir / "context.txt"
        context = context_file.read_text(encoding="utf-8").strip() if context_file.is_file() else None
        messages = self._build_messages(service_dir.name, examples, release_notes, current, latest, context)

        try:
            return self._classify_with_model(self.model, messages, timeout)
        except Exception as primary_exc:
            if not self.fallback_model or self.fallback_model == self.model:
                return {"action": "manual", "reason": f"Classification error: {primary_exc}"}
            try:
                return self._classify_with_model(self.fallback_model, messages, timeout)
            except Exception as fallback_exc:
                return {
                    "action": "manual",
                    "reason": f"Classification error: {primary_exc}; fallback ({self.fallback_model}) also failed: {fallback_exc}",
                }

    def _classify_with_model(self, model: str, messages: List[Dict[str, str]], timeout: float) -> Dict[str, str]:
        try:
            payload = json.dumps({
                "model": model,
                "messages": messages,
                "stream": False,
                "options": {"temperature": 0},
            }).encode("utf-8")
            headers = {"Content-Type": "application/json"}
            if self.api_key:
                headers["Authorization"] = f"Bearer {self.api_key}"
            req = Request(
                f"{self.base_url}/v1/chat/completions",
                data=payload,
                headers=headers,
                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:
            raise RuntimeError(f"model {model}: {exc}") from exc

    def _build_messages(
        self,
        service_name: str,
        examples: List[Dict[str, str]],
        release_notes: str,
        current: str,
        latest: str,
        context: Optional[str] = None,
    ) -> 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"
        )
        if context:
            system_prompt += f"Instance context: {context}\n\n"
        system_prompt += (
            "Rules:\n"
            "- Only flag as manual if the breaking changes actually affect this specific setup.\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

    @staticmethod
    def _truncate(text: str, max_chars: int = 16000) -> str:
        if len(text) <= max_chars:
            return text

        # Split into sections by markdown headers
        import re
        sections = re.split(r'(?=^#{1,4}\s)', text, flags=re.MULTILINE)

        important_keywords = re.compile(
            r'breaking|deprecat|migrat|important|caution|warning|manual|required|must|before.*upgrad',
            re.IGNORECASE,
        )

        # Always keep the intro (first section) and any important sections
        kept = []
        skipped = 0
        for i, section in enumerate(sections):
            if i == 0 or important_keywords.search(section[:500]):
                kept.append(section)
            else:
                skipped += 1

        result = "".join(kept)
        if skipped:
            result += f"\n\n[... {skipped} sections of features/bug fixes omitted ...]"

        # Final safety truncation if still too long
        if len(result) > max_chars:
            result = result[:max_chars] + "\n\n[... truncated ...]"

        return result