diff --git a/test_classifier.py b/test_classifier.py index 93cc9ce..5b50590 100644 --- a/test_classifier.py +++ b/test_classifier.py @@ -118,10 +118,13 @@ def call_ollama(messages: list[dict]) -> dict: "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={"Content-Type": "application/json"}, + headers=headers, method="POST", ) with urlopen(req, timeout=300) as resp: diff --git a/update_monitor/classifier.py b/update_monitor/classifier.py index 9a072d1..038d0f4 100644 --- a/update_monitor/classifier.py +++ b/update_monitor/classifier.py @@ -34,6 +34,7 @@ class Classifier: 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, @@ -77,10 +78,13 @@ class Classifier: "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={"Content-Type": "application/json"}, + headers=headers, method="POST", ) with urlopen(req, timeout=timeout) as resp: