checks/health_check.py

# checks/health_check.py
import os
import requests
from typing import Dict, Tuple

from checks.endpoint_check import check_public_endpoint

def check_with_prompt(endpoint_uri: str):
    
    # Try a direct test of your endpoint
    test_prompt = "Hello, how are you?"
    response = requests.post(
        endpoint_uri,
        json={"inputs": test_prompt},
        headers={"Authorization": f"Bearer {os.getenv('HF_TOKEN')}"}
    )
    print(response.json())



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