Add new NIM models + background health check loop

backend.py

@@ -27,6 +27,7 @@ from fastapi import FastAPI, Request, Header, HTTPException
 from fastapi.responses import StreamingResponse, JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 from openai import AsyncOpenAI
+import anyio
 import asyncpg
 
 # ---------------------------------------------------------------------------
@@ -41,6 +42,12 @@ MAX_TOOL_ROUNDS = int(os.environ.get("MAX_TOOL_ROUNDS", "10"))
 
 # NIM models that reliably support tool/function calling
 TOOL_CAPABLE_MODELS = {
+    "nvidia/nemotron-3-ultra-550b-a55b": "Nemotron 3 Ultra 550B (Agentic)",
+    "z-ai/glm-5.1": "GLM 5.1 (Agentic)",
+    "moonshotai/kimi-k2.6": "Kimi K2.6 (Agentic)",
+    "minimaxai/minimax-m3": "MiniMax M3 (Agentic)",
+    "stepfun-ai/step-3.7-flash": "Step 3.7 Flash (Agentic)",
+    "minimaxai/minimax-m2.7": "MiniMax M2.7 (Agentic)",
     "meta/llama-3.1-70b-instruct": "Llama 3.1 70B (Agentic)",
     "meta/llama-3.1-405b-instruct": "Llama 3.1 405B (Agentic)",
     "qwen/qwen2.5-coder-32b-instruct": "Qwen 2.5 Coder 32B (Agentic)",
@@ -55,6 +62,8 @@ ALL_MODELS = {
     "mistralai/mistral-large-2-instruct": "Mistral Large 2 (Chat only)",
 }
 
+RECOMMENDED_MODEL = "nvidia/llama-3.1-nemotron-70b-instruct"
+
 # Ensure workspace exists
 Path(WORKSPACE_DIR).mkdir(parents=True, exist_ok=True)
 
@@ -400,10 +409,70 @@ def auth(authorization: str = None):
         raise HTTPException(status_code=401, detail="Unauthorized")
 
 
+async def check_models_health():
+    global RECOMMENDED_MODEL
+    # Top models to test
+    models_to_test = [
+        "nvidia/nemotron-3-ultra-550b-a55b",
+        "z-ai/glm-5.1",
+        "moonshotai/kimi-k2.6",
+        "minimaxai/minimax-m3",
+        "stepfun-ai/step-3.7-flash",
+        "minimaxai/minimax-m2.7",
+        "nvidia/llama-3.1-nemotron-70b-instruct",
+        "meta/llama-3.1-405b-instruct",
+        "qwen/qwen2.5-coder-32b-instruct",
+        "meta/llama-3.3-70b-instruct",
+    ]
+    
+    best_model = None
+    best_latency = 999.0
+    
+    print("[Health Check] Starting periodic model verification...")
+    for model in models_to_test:
+        start_time = time.time()
+        try:
+            # Send a fast test prompt
+            async with anyio.fail_after(15.0): # 15 seconds max timeout
+                await nim_client.chat.completions.create(
+                    model=model,
+                    messages=[{"role": "user", "content": "1+1="}],
+                    max_tokens=3,
+                )
+            latency = time.time() - start_time
+            print(f"[Health Check] Model {model} is ONLINE. Latency: {latency:.2f}s")
+            
+            # We want the model that is within 15 seconds
+            # and is the fastest (lowest latency)
+            if latency < best_latency:
+                best_latency = latency
+                best_model = model
+                
+        except Exception as e:
+            print(f"[Health Check] Model {model} is OFFLINE or TIMEOUT: {e}")
+            
+    if best_model:
+        RECOMMENDED_MODEL = best_model
+        print(f"[Health Check] Best model found: {RECOMMENDED_MODEL} ({best_latency:.2f}s)")
+    else:
+        print("[Health Check] Warning: All checked models failed or timed out!")
+
+async def periodic_health_check_loop():
+    # Wait 30 seconds after startup before the first check to let the space boot fully
+    await asyncio.sleep(30)
+    while True:
+        try:
+            await check_models_health()
+        except Exception as e:
+            print(f"[Health Check] Loop error: {e}")
+        await asyncio.sleep(300) # every 5 minutes
+
 @app.on_event("startup")
 async def startup():
     await init_db()
     Path(WORKSPACE_DIR).mkdir(parents=True, exist_ok=True)
+    # Start background health checking
+    asyncio.create_task(periodic_health_check_loop())
     print(f"[Backend] Started. Workspace: {WORKSPACE_DIR}")
     print(f"[Backend] Tool-capable models: {list(TOOL_CAPABLE_MODELS.keys())}")
     print(f"[Backend] DB connected: {db_pool is not None}")
@@ -656,6 +725,7 @@ async def health():
         "workspace_exists": Path(WORKSPACE_DIR).exists(),
         "db_connected": db_pool is not None,
         "models_count": len(ALL_MODELS),
+        "recommended_model": RECOMMENDED_MODEL,
     }