Update endpoint_utils.py

endpoint_utils.py

@@ -4,7 +4,10 @@ from typing import Optional, Tuple, Callable, Dict, Any
 from urllib.parse import urlparse
 import os, time, requests
 
-
+def hf_headers():
+    tok = os.environ.get("HF_TOKEN", "").strip()
+    return {"Authorization": f"Bearer {tok}"} if tok else {}
+    
 def _valid_uri(uri: Optional[str]) -> bool:
     if not uri:
         return False
@@ -12,6 +15,13 @@ def _valid_uri(uri: Optional[str]) -> bool:
     return p.scheme in {"http", "https"} and bool(p.netloc)
 
 
+def _detail(resp: requests.Response) -> str:
+    try:
+        j = resp.json()
+        return (j.get("error") or j.get("message") or "").strip()
+    except Exception:
+        return (resp.text or "").strip()
+
 def wake_endpoint(
     uri: Optional[str],
     *,
@@ -28,35 +38,34 @@ def wake_endpoint(
     if not _valid_uri(uri):
         return False, "invalid or missing URI (expect http(s)://...)"
 
-    headers: Dict[str, str] = {}
-    tok = token or os.environ.get("HF_TOKEN")
-    if tok:
-        headers["Authorization"] = f"Bearer {tok}"
+    headers = hf_headers()
+    if not headers:
+        log("⚠️ HF_TOKEN not set — POST / will likely return 401/403.")
 
-    # 0) Try a quick health check first (cheap)
-    last_detail = "no response"
+    last = "no response"
+    
+     # /health probe (auth included if required)
     try:
         hr = requests.get(f"{uri.rstrip('/')}/health", headers=headers, timeout=5)
         if hr.ok:
             log("✅ /health reports ready.")
             return True, None
-        try:
-            last_detail = (hr.json().get("error") or hr.json().get("message"))  # type: ignore
-        except Exception:
-            last_detail = (hr.text or "").strip()
-        log(f"[health] HTTP {hr.status_code} – {last_detail or 'warming?'}")
+        last = f"HTTP {hr.status_code} – {_detail(hr) or 'warming?'}"
+        log(f"[health] {last}")
+        if hr.status_code in (401, 403):
+            return False, f"Unauthorized (check HF_TOKEN). {last}"
     except requests.RequestException as e:
-        last_detail = type(e).__name__
-        log(f"[health] {last_detail}")
+        last = type(e).__name__
+        log(f"[health] {last}")
 
-    # 1) Initial nudge (ignore errors)
+    # warmup nudge
     payload = warm_payload if warm_payload is not None else {"inputs": "wake"}
     try:
         requests.post(uri, headers=headers, json=payload, timeout=5)
     except requests.RequestException:
         pass
 
-    # 2) Poll until healthy or timeout
+    # poll
     deadline = time.time() + max_wait
     while time.time() < deadline:
         try:
@@ -65,24 +74,22 @@ def wake_endpoint(
                 log("✅ Endpoint is awake and responsive.")
                 return True, None
 
-            # extract any helpful server message
-            detail = ""
-            try:
-                data = r.json()
-                detail = data.get("error") or data.get("message") or ""
-            except ValueError:
-                detail = (r.text or "").strip()
+            d = _detail(r)
+            last = f"HTTP {r.status_code}" + (f" – {d}" if d else "")
+
+            if r.status_code in (401, 403):
+                return False, f"Unauthorized (check HF_TOKEN, org access). {last}"
 
-            last_detail = f"HTTP {r.status_code}" + (f" – {detail}" if detail else "")
             if r.status_code in (429, 503, 504):
-                log(f"[server] {detail or 'warming up'} (HTTP {r.status_code}); retrying in {int(poll_every)}s…")
+                log(f"[server] {d or 'warming up'} (HTTP {r.status_code}); retrying in {int(poll_every)}s…")
             else:
-                log(f"[server] {detail or 'unexpected response'} (HTTP {r.status_code}); retrying in {int(poll_every)}s…")
+                log(f"[server] {d or 'unexpected response'} (HTTP {r.status_code}); retrying in {int(poll_every)}s…")
 
         except requests.RequestException as e:
-            last_detail = type(e).__name__
-            log(f"[client] {last_detail}; retrying in {int(poll_every)}s…")
+            last = type(e).__name__
+            log(f"[client] {last}; retrying in {int(poll_every)}s…")
 
         time.sleep(poll_every)
 
-    return False, f"Timed out after {max_wait}s — last status: {last_detail}"
+    return False, f"Timed out after {max_wait}s — last status: {last}"
+