Create health_check.py

checks/health_check.py

@@ -0,0 +1,38 @@
+# checks/health_check.py
+import os
+import requests
+from typing import Dict, Tuple
+
+def check_model_endpoint(endpoint_uri: str) -> Tuple[bool, Dict]:
+    """Comprehensive check of model endpoint health"""
+    # Basic availability check
+    status_info = check_public_endpoint(endpoint_uri)
+    
+    if not status_info["status"] or status_info["status_code"] in [503, 502]:
+        return False, status_info
+    
+    # Test actual model response
+    try:
+        test_prompt = "Hello, how are you?"
+        response = requests.post(
+            endpoint_uri,
+            json={"inputs": test_prompt},
+            headers={"Authorization": f"Bearer {os.getenv('HF_TOKEN')}"},
+            timeout=10
+        )
+        response.raise_for_status()
+        return True, {
+            "status": True,
+            "status_code": response.status_code,
+            "response": response.json()
+        }
+    except Exception as e:
+        return False, {
+            "status": False,
+            "status_code": 500,
+            "error": str(e)
+        }
+
+def should_launch_ui(status_info: Dict) -> bool:
+    """Determine if UI should be launched based on status"""
+    return status_info.get("status", False) and status_info.get("status_code", 500) == 200