my-homelab-configs/scripts/ai-model-check.py

16 lines
500 B
Python

import json
import sys
import urllib.error
import urllib.request
endpoint, model = sys.argv[1:3]
try:
with urllib.request.urlopen(
"{}/api/tags".format(endpoint.rstrip("/")), timeout=5
) as response:
tags = json.loads(response.read().decode("utf-8"))
models = {item.get("name", "") for item in tags.get("models", [])}
sys.exit(0 if model in models or f"{model}:latest" in models else 1)
except (urllib.error.URLError, json.JSONDecodeError, KeyError):
sys.exit(1)