import litellm
import asyncio

async def check_model_health(model_name):
    try:
        response = await litellm.acompletion(
            model=model_name,
            messages=[{"role": "user", "content": "Are you healthy?"}],
            timeout=10  # Add timeout to prevent indefinite hanging
        )
        if response and response.choices and len(response.choices) > 0:
            return {"model": model_name, "status": "healthy", "response": response.choices[0].message.content}
        else:
            return {"model": model_name, "status": "unhealthy", "error": "No response content"}
    except Exception as e:
        return {"model": model_name, "status": "unhealthy", "error": str(e)}

async def get_all_model_health(model_list):
    tasks = [check_model_health(model) for model in model_list]
    results = await asyncio.gather(*tasks)
    return results


async def main():
    # Replace with your actual model list from the router configuration
    model_list = litellm.model_list # Assuming litellm has access to the model list
    if not model_list:
        print("Error: Model list is empty or unavailable.")
        return
    health_results = await get_all_model_health(model_list)
    for result in health_results:
        print(f"Model: {result['model']}, Status: {result['status']}")
        if result['status'] == 'unhealthy':
            print(f"  Error: {result['error']}")

if __name__ == "__main__":
    asyncio.run(main())