#!/usr/bin/env python3 """ Generate LiteLLM config from verified working NVIDIA NIM models. Tests at startup — only confirmed-working models are included. """ import json, os, requests, yaml, sys, concurrent.futures NVIDIA_API_BASE = "https://integrate.api.nvidia.com/v1" WORKING_MODELS = [ "deepseek-ai/deepseek-v4-pro", "google/gemma-2-2b-it", "google/gemma-3n-e2b-it", "google/gemma-3n-e4b-it", "meta/llama-3.1-70b-instruct", "meta/llama-3.1-8b-instruct", "meta/llama-3.2-11b-vision-instruct", "meta/llama-3.2-1b-instruct", "meta/llama-3.2-3b-instruct", "meta/llama-3.2-90b-vision-instruct", "meta/llama-3.3-70b-instruct", "meta/llama-4-maverick-17b-128e-instruct", "mistralai/ministral-14b-instruct-2512", "mistralai/mistral-7b-instruct-v0.3", "mistralai/mistral-large-3-675b-instruct-2512", "mistralai/mistral-nemotron", "mistralai/mistral-small-4-119b-2603", "mistralai/mixtral-8x7b-instruct-v0.1", "nvidia/gliner-pii", "nvidia/ising-calibration-1-35b-a3b", "nvidia/llama-3.1-nemotron-nano-8b-v1", "nvidia/llama-3.1-nemotron-nano-vl-8b-v1", "nvidia/llama-3.3-nemotron-super-49b-v1", "nvidia/llama-3.3-nemotron-super-49b-v1.5", "nvidia/nemotron-3-nano-30b-a3b", "nvidia/nemotron-3-nano-omni-30b-a3b-reasoning", "nvidia/nemotron-3-super-120b-a12b", "nvidia/nemotron-mini-4b-instruct", "nvidia/nemotron-nano-12b-v2-vl", "nvidia/nvidia-nemotron-nano-9b-v2", "openai/gpt-oss-120b", "openai/gpt-oss-20b", "qwen/qwen3-coder-480b-a35b-instruct", "qwen/qwen3-next-80b-a3b-instruct", "qwen/qwen3.5-122b-a10b", "qwen/qwen3.5-397b-a17b", "sarvamai/sarvam-m", "stepfun-ai/step-3.5-flash", "stepfun-ai/step-3.7-flash", "stockmark/stockmark-2-100b-instruct", "upstage/solar-10.7b-instruct", ] def test_model(mid: str, api_key: str, timeout: int = 20) -> bool: try: r = requests.post(f"{NVIDIA_API_BASE}/chat/completions", headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}, json={"model": mid, "messages": [{"role": "user", "content": "hi"}], "max_tokens": 3}, timeout=timeout) return r.status_code == 200 except: return False def get_tier(mid: str) -> str: m = mid.lower() if "coder" in m or "codestral" in m or "deepseek-v4" in m: return "coding" if "qwen3.5-397b" in m or "mistral-large-3" in m or "nemotron-3-super" in m or "gpt-oss-120b" in m or "deepseek-v4" in m: return "reasoning" if "llama-3.3-70b" in m or "llama-4-maverick" in m or "mistral-nemotron" in m or "stockmark" in m or "mixtral" in m: return "general" if "llama-3.1-70b" in m or "mistral-small" in m or "qwen3-next" in m or "gemma-3n" in m or "nemotron-3-nano" in m or "nvidia-llama-3.3-nemotron" in m: return "general" if "qwen3.5-122b" in m or "ministral" in m or "llama-3.1-nemotron-nano" in m or "gpt-oss-20b" in m: return "reasoning" if "llama-3.2-90b-vision" in m or "llama-3.2-11b-vision" in m or "nemotron-nano-12b-v2-vl" in m: return "vision" if "solar" in m or "step" in m or "sarvam" in m: return "reasoning" return "fast" def main(): api_key = os.environ.get("NVIDIA_API_KEY", "") if not api_key: print("ERROR: NVIDIA_API_KEY not set") sys.exit(1) print("[*] Testing working models against NVIDIA NIM API...") candidates = list(WORKING_MODELS) working = [] with concurrent.futures.ThreadPoolExecutor(max_workers=8) as pool: fut_to_mid = {pool.submit(test_model, m, api_key): m for m in candidates} for fut in concurrent.futures.as_completed(fut_to_mid): mid = fut_to_mid[fut] if fut.result(): working.append(mid) print(f" OK: {mid}") else: print(f" SKIP: {mid}") if not working: print("[!] All models failed — using hardcoded list anyway") working = list(WORKING_MODELS) working.sort() print(f"\n[+] {len(working)} working models") tier_models = {} for m in working: t = get_tier(m) tier_models.setdefault(t, []).append(m) model_list = [] for mid in working: name = mid.split("/")[-1].replace("-instruct", "").replace("-v0.1", "").replace("-2512", "").replace("-2603", "") model_list.append({"id": mid, "name": name, "tier": get_tier(mid)}) litellm_list = [] short_names = set() for m in model_list: litellm_list.append({ "model_name": "nvidia-auto", "litellm_params": { "model": f"nvidia_nim/{m['id']}", "api_key": "os.environ/NVIDIA_API_KEY", "api_base": NVIDIA_API_BASE, "timeout": 30, "stream_timeout": 60, }, }) for t, mids in tier_models.items(): for mid in mids: litellm_list.append({ "model_name": f"nvidia-{t}", "litellm_params": { "model": f"nvidia_nim/{mid}", "api_key": "os.environ/NVIDIA_API_KEY", "api_base": NVIDIA_API_BASE, "timeout": 30, "stream_timeout": 60, }, }) for m in model_list: short = m["id"].split("/")[-1] if short not in short_names: short_names.add(short) litellm_list.append({ "model_name": short, "litellm_params": { "model": f"nvidia_nim/{m['id']}", "api_key": "os.environ/NVIDIA_API_KEY", "api_base": NVIDIA_API_BASE, }, }) fallbacks = [ {"nvidia-coding": ["nvidia-reasoning", "nvidia-general"]}, {"nvidia-reasoning": ["nvidia-general", "nvidia-fast"]}, {"nvidia-general": ["nvidia-reasoning", "nvidia-fast"]}, {"nvidia-vision": ["nvidia-reasoning", "nvidia-general"]}, {"nvidia-fast": ["nvidia-general"]}, {"nvidia-auto": ["nvidia-reasoning", "nvidia-general"]}, ] config = { "model_list": litellm_list, "litellm_settings": { "num_retries": 3, "request_timeout": 30, "fallbacks": fallbacks, "set_verbose": False, "drop_params": True, }, "router_settings": { "routing_strategy": "latency-based-routing", "num_retries": 3, "cooldown_time": 60, "retry_after": 5, "allowed_fails": 2, }, "general_settings": { "master_key": "sk-litellm-master", }, } with open("config.yaml", "w") as f: yaml.dump(config, f, default_flow_style=False, sort_keys=False, width=120) with open("models.json", "w") as f: json.dump(model_list, f, indent=2) print(f"\n[+] Config: {len(working)} models, {len(tier_models)} tiers") for t, mids in tier_models.items(): print(f" {t}: {', '.join(m.split('/')[-1][:35] for m in mids)}") if __name__ == "__main__": main()