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.
200 lines
6.9 KiB
Python
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()
|