Spaces:
Sleeping
Sleeping
| #!/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() | |