service-update-alerts/test_classifier.py

#!/usr/bin/env python3
"""
Test the Ollama model's ability to classify Immich release notes.
Runs leave-one-out: for each example, uses the rest as few-shot context
and checks whether the model predicts the correct action.
"""

import json
import os
import sys
from pathlib import Path
from urllib.request import Request, urlopen

OLLAMA_BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://192.168.1.20:11434")
OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "gemma4:26b-64k")
EXAMPLES_DIR = Path(__file__).parent / "services" / "immich" / "review-examples"


def load_examples(directory: Path) -> list[dict]:
    examples = []
    for path in sorted(directory.glob("*.txt")):
        text = path.read_text(encoding="utf-8")
        if "---" not in text:
            print(f"WARNING: {path.name} has no frontmatter separator, skipping")
            continue
        header, _, body = text.partition("---\n")
        action = None
        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"):
            print(f"WARNING: {path.name} has invalid action '{action}', skipping")
            continue
        examples.append({
            "file": path.name,
            "action": action,
            "reason": reason or "",
            "notes": body.strip(),
        })
    return examples


def build_messages(examples: list[dict], target_notes: str) -> list[dict]:
    system = (
        "You are an update safety classifier for self-hosted Immich instances.\n"
        "Given release notes for an Immich update, decide whether the update can be\n"
        "applied automatically (action: auto) or requires the admin to intervene\n"
        "manually before or 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"
        "- Respond ONLY with a JSON object: {\"action\": \"auto\", \"reason\": \"...\"}\n"
        "  or {\"action\": \"manual\", \"reason\": \"...\"}\n"
        "- Keep reason under 60 words."
    )

    messages = [{"role": "system", "content": system}]

    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"Release notes:\n{target_notes}"})
    return messages


def call_ollama(messages: list[dict]) -> dict:
    payload = json.dumps({
        "model": OLLAMA_MODEL,
        "messages": messages,
        "stream": False,
        "options": {"temperature": 0},
    }).encode("utf-8")
    req = Request(
        f"{OLLAMA_BASE_URL}/v1/chat/completions",
        data=payload,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urlopen(req, timeout=300) as resp:
        data = json.loads(resp.read())
    return data


def parse_response(text: str) -> dict:
    text = text.strip()
    # extract JSON from response (model may wrap it in markdown)
    start = text.find("{")
    end = text.rfind("}") + 1
    if start == -1 or end == 0:
        return {"action": "parse_error", "reason": text[:200]}
    try:
        return json.loads(text[start:end])
    except json.JSONDecodeError:
        return {"action": "parse_error", "reason": text[:200]}


def main():
    examples = load_examples(EXAMPLES_DIR)
    if not examples:
        print("No examples found.")
        sys.exit(1)

    print(f"Loaded {len(examples)} examples from {EXAMPLES_DIR}")
    print(f"Model: {OLLAMA_MODEL}")
    print(f"Running leave-one-out test ({len(examples)} calls)...\n")

    passed = 0
    failed = 0

    for i, target in enumerate(examples):
        few_shot = [ex for j, ex in enumerate(examples) if j != i]
        messages = build_messages(few_shot, target["notes"])

        try:
            response = call_ollama(messages)
            content = response["choices"][0]["message"]["content"]
            result = parse_response(content)
        except Exception as exc:
            result = {"action": "error", "reason": str(exc)}

        predicted = result.get("action", "error")
        expected = target["action"]
        ok = predicted == expected
        status = "PASS" if ok else "FAIL"

        if ok:
            passed += 1
        else:
            failed += 1

        print(f"[{status}] {target['file']}")
        print(f"       expected : {expected}")
        print(f"       predicted: {predicted}")
        print(f"       reason   : {result.get('reason', '')}")
        print()

    print(f"Results: {passed}/{len(examples)} passed", end="")
    if failed:
        print(f", {failed} failed")
    else:
        print(" — all correct!")

    sys.exit(0 if failed == 0 else 1)


if __name__ == "__main__":
    main()