service-update-alerts/test_classifier.py
Jeena 3cba601c72 classifier: Support authenticated Ollama endpoints via OLLAMA_API_KEY
Ollama Cloud and other hosted endpoints reject unauthenticated requests,
so the classifier and its test harness could not reach them. Send a
Bearer token when OLLAMA_API_KEY is set, letting the same code work
against both the local instance and hosted services without changing
behavior for keyless setups.
2026-10-08 08:26:26 +09:00

200 lines
6.9 KiB
Python

#!/usr/bin/env python3
"""
Test the Ollama model's ability to classify 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.
Usage: python test_classifier.py [service_name]
service_name defaults to "immich"
"""
import json
import os
import sys
from pathlib import Path
from urllib.request import Request, urlopen
import re as _re
OLLAMA_BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://192.168.1.20:11434")
OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "gemma4:26b")
SERVICE_NAME = sys.argv[1] if len(sys.argv) > 1 else "immich"
EXAMPLES_DIR = Path(__file__).parent / "services" / SERVICE_NAME / "review-examples"
def guess_version_from_notes(notes: str) -> str:
"""Extract a version number from release notes (e.g. '# v2.7.0' or '### v8.1.5')."""
match = _re.search(r"v?(\d+\.\d+\.\d+)", notes)
return match.group(1) if match else "1.0.0"
def make_previous_version(version: str) -> str:
"""Decrement version to simulate upgrading from a previous release."""
parts = version.split(".")
if len(parts) == 3 and parts[2].isdigit():
patch = int(parts[2])
if patch > 0:
return f"{parts[0]}.{parts[1]}.{patch - 1}"
# x.y.0 → decrement minor
minor = int(parts[1])
if minor > 0:
return f"{parts[0]}.{minor - 1}.0"
# x.0.0 → decrement major
major = int(parts[0])
if major > 0:
return f"{major - 1}.0.0"
return version
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, current: str = "1.0.0", latest: str = "2.0.0") -> list[dict]:
context_file = EXAMPLES_DIR.parent / "context.txt"
context = context_file.read_text(encoding="utf-8").strip() if context_file.is_file() else None
system = (
f"You are an update safety classifier for self-hosted {SERVICE_NAME} instances.\n"
f"Given release notes for a {SERVICE_NAME} 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"
)
if context:
system += f"Instance context: {context}\n\n"
system += (
"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 = [{"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"Upgrading from {current} to {latest}.\n\nRelease 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")
headers = {"Content-Type": "application/json"}
if api_key := os.getenv("OLLAMA_API_KEY", ""):
headers["Authorization"] = f"Bearer {api_key}"
req = Request(
f"{OLLAMA_BASE_URL}/v1/chat/completions",
data=payload,
headers=headers,
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]
latest = guess_version_from_notes(target["notes"])
current = make_previous_version(latest)
messages = build_messages(few_shot, target["notes"], current, latest)
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()