NIM Model Router

OpenAI-compatible proxy that routes requests to the best NVIDIA NIM model by task.

NVIDIA's NIM catalog has 100+ models. Picking the right one for each request is tedious. This router sits in front of the NIM API and automatically selects a model based on what you're asking for — fast chat, agentic tool use, deep reasoning, coding, embeddings, reranking, and more.

Drop it into any OpenAI SDK client by changing base_url. No other code changes required.

Quick start

git clone https://github.com/cobusgreyling/nim-model-router.git
cd nim-model-router
python -m venv .venv && source .venv/bin/activate
pip install -e ".[dev]"

cp .env.example .env
# Edit .env and set NVIDIA_API_KEY

nim-router serve

The proxy listens on http://127.0.0.1:8080. API docs: http://127.0.0.1:8080/docs

Docker

cp .env.example .env  # set NVIDIA_API_KEY
docker compose up --build

Usage

Auto-routing (recommended)

from openai import OpenAI

client = OpenAI(
    base_url="http://127.0.0.1:8080/v1",
    api_key="local",  # not used — router injects NVIDIA_API_KEY upstream
)

response = client.chat.completions.create(
    model="nim-router/auto",
    messages=[{"role": "user", "content": "Build a Python agent with tool calling"}],
)
print(response.choices[0].message.content)

Explicit task aliases

Force a task via header

curl http://127.0.0.1:8080/v1/chat/completions \
  -H "Content-Type: application/json" \
  -H "X-NIM-Task: reasoning" \
  -d '{
    "model": "nim-router/auto",
    "messages": [{"role": "user", "content": "Analyze the root cause step by step"}]
  }'

Rerank

curl http://127.0.0.1:8080/v1/rerank \
  -H "Content-Type: application/json" \
  -d '{
    "model": "nim-router/rerank",
    "query": "What is NVIDIA NIM?",
    "documents": ["NIM provides optimized inference...", "Unrelated text"],
    "top_n": 2
  }'

Passthrough to a specific NIM model

If you pass a concrete NIM model ID (e.g. meta/llama-3.1-70b-instruct), the router forwards it unchanged.

How routing works

flowchart TD
    A[Client request] --> B{Model alias?}
    B -->|nim-router/*| C[Resolve task]
    B -->|concrete NIM ID| D[Passthrough]
    B -->|nim-router/auto| E[Classifier]
    A --> F{X-NIM-Task header?}
    F -->|set| C
    E --> G{Signals}
    G -->|tools| H[agentic]
    G -->|large prompt| I[long_context]
    G -->|keywords| J[coding / reasoning / rerank]
    G -->|short prompt| K[fast]
    G -->|ambiguous| L[general]
    C --> M[Apply policies]
    M --> N[Pick model + fallbacks]
    N --> O[Proxy to NIM API]
    O -->|5xx| P[Try fallback model]

Classifier signals:

• Tools present → agentic
• Large prompt (>12k estimated tokens, tiktoken) → long_context
• Reasoning keywords → reasoning
• Coding keywords → coding
• Rerank keywords / query+documents → rerank
• Short prompt (≤120 chars) → fast
• Ambiguous → general (not ultra-expensive agentic)

Policies can downgrade ultra models for short prompts and route low-confidence requests to general.

CLI

# Start proxy
nim-router serve --port 8080 --config src/nim_model_router/models.yaml

# Dry-run routing (no API call)
nim-router route "refactor this Python function" --json
nim-router route "hello"
nim-router route "plan a multi-step agent" --tools

# Show registry
nim-router models

# Sync model suggestions from NIM catalog
nim-router catalog-sync --task coding

# Print OpenAI SDK example
nim-router client-example

Observability

Every proxied response includes routing metadata:

# Live stats
curl http://127.0.0.1:8080/v1/router/stats

# Task registry + cost summary
curl http://127.0.0.1:8080/v1/router/tasks

# Reload config without restart
curl -X POST http://127.0.0.1:8080/v1/router/reload

# Prometheus metrics
curl http://127.0.0.1:8080/metrics

# Dry-run endpoint
curl -X POST http://127.0.0.1:8080/v1/router/dry-run \
  -H "Content-Type: application/json" \
  -d '{"messages":[{"role":"user","content":"debug my rust code"}]}'

Set ROUTER_LOG_PATH=data/router.log.jsonl to persist request logs.

Environment variables

Customizing models

Edit src/nim_model_router/models.yaml (or set ROUTER_CONFIG):

tasks:
  agentic:
    model: nvidia/nemotron-3-nano-30b-a3b
    fallbacks:
      - general
      - fast
    extra_body:
      enable_thinking: true
      reasoning_budget: 2048
    ab_test:
      enabled: false
      variants:
        - model: nvidia/nemotron-3-nano-30b-a3b
          weight: 50
        - model: nvidia/nemotron-3-super-120b-a12b
          weight: 50

classifier:
  use_llm_classifier: false  # set true + pip install ".[llm-classifier]"

Reload at runtime: curl -X POST http://127.0.0.1:8080/v1/router/reload

Integrations

LiteLLM

model_list:
  - model_name: nim-auto
    litellm_params:
      model: openai/nim-router/auto
      api_base: http://127.0.0.1:8080/v1
      api_key: local

LangChain

from langchain_openai import ChatOpenAI

llm = ChatOpenAI(
    base_url="http://127.0.0.1:8080/v1",
    api_key="local",
    model="nim-router/auto",
)

Continue / Cursor (OpenAI-compatible)

{
  "models": [{
    "title": "NIM Auto",
    "provider": "openai",
    "model": "nim-router/auto",
    "apiBase": "http://127.0.0.1:8080/v1",
    "apiKey": "local"
  }]
}

Development

pip install -e ".[dev]"
pytest --cov=nim_model_router --cov-report=term-missing
ruff check src tests
ruff format src tests

See CONTRIBUTING.md.

Security

• Store NVIDIA_API_KEY in .env — never commit it.
• Set ROUTER_API_KEY before exposing the proxy beyond localhost.
• Bind to 127.0.0.1 by default. Use Docker/reverse proxy auth for production.

License

MIT