Building a router that picks the best model per task type

A router best model per task lets you match each request to the model that balances cost, latency, and quality. This tutorial builds a working Python router that classifies incoming prompts and dispatches them to a curated model map, with fallback and cache control wired in from the start.

Prerequisites

• Python 3.11+ with openai>=1.40 installed
• An API key for an OpenAI-compatible inference gateway
• Familiarity with async Python and JSON schemas

We’ll point the client at a single OpenAI-compatible endpoint (n4n.ai addresses 240+ models behind one URL and honors client routing directives). The routing logic is provider-agnostic; swap the base URL if you self-host.

pip install openai>=1.40
export LLM_API_KEY=sk-...
export LLM_BASE_URL=https://api.n4n.ai/v1

1. Define task types and model map

Start with an explicit contract. We support four task classes: code, math, summarize, chat. Each maps to a primary model and a cheaper fallback.

from enum import Enum
from dataclasses import dataclass

class TaskType(str, Enum):
    CODE = "code"
    MATH = "math"
    SUMMARIZE = "summarize"
    CHAT = "chat"

@dataclass
class ModelRoute:
    primary: str
    fallback: str
    max_tokens: int

ROUTES: dict[TaskType, ModelRoute] = {
    TaskType.CODE: ModelRoute("meta-llama/codellama-34b-instruct", "openai/gpt-4o-mini", 1024),
    TaskType.MATH: ModelRoute("anthropic/claude-3-5-sonnet", "openai/gpt-4o", 512),
    TaskType.SUMMARIZE: ModelRoute("google/gemini-flash-1.5", "openai/gpt-4o-mini", 256),
    TaskType.CHAT: ModelRoute("openai/gpt-4o-mini", "meta-llama/llama-3.1-8b-instruct", 512),
}

Keep this table in configuration, not code, in production. Hardcoding is fine for a tutorial.

2. Build a task classifier

Calling a heavy model to classify wastes tokens. Use a small model with a strict JSON schema. We’ll use gpt-4o-mini as the classifier.

from openai import OpenAI
import os, json

client = OpenAI(
    api_key=os.environ["LLM_API_KEY"],
    base_url=os.environ["LLM_BASE_URL"],
)

CLASSIFIER_MODEL = "openai/gpt-4o-mini"

SYSTEM_PROMPT = """Classify the user request into exactly one of:
code, math, summarize, chat.
Respond with JSON: {"task": "<type>"}"""

def classify(task_text: str) -> TaskType:
    resp = client.chat.completions.create(
        model=CLASSIFIER_MODEL,
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": task_text},
        ],
        response_format={"type": "json_object"},
        max_tokens=32,
    )
    raw = json.loads(resp.choices[0].message.content)
    return TaskType(raw["task"])

Test it:

print(classify("Write a bash script to rotate logs"))
# TaskType.CODE
print(classify("What is 17 * 23?"))
# TaskType.MATH

Expected output at checkpoint:

TaskType.CODE
TaskType.MATH

If the classifier returns an unknown label, default to CHAT.

3. Dispatch with routing directives

The router best model per task must forward the chosen model and respect cache hints. We’ll pass extra_headers for provider cache control where supported.

from typing import Optional

def complete(task_text: str, task_type: Optional[TaskType] = None) -> str:
    task_type = task_type or classify(task_text)
    route = ROUTES[task_type]

    try:
        resp = client.chat.completions.create(
            model=route.primary,
            messages=[{"role": "user", "content": task_text}],
            max_tokens=route.max_tokens,
            extra_headers={"x-cache-control": "ttl=300"},  # hint gateway to cache
        )
        return resp.choices[0].message.content
    except Exception as e:
        # automatic fallback on rate limit / degradation
        resp = client.chat.completions.create(
            model=route.fallback,
            messages=[{"role": "user", "content": task_text}],
            max_tokens=route.max_tokens,
        )
        return resp.choices[0].message.content

The x-cache-control header is a client routing directive that compliant gateways forward to the upstream provider. If the primary model is rate-limited, we catch and hit the fallback.

4. Add a lightweight cache layer

Classified routing pairs well with a local response cache for repeated tasks (e.g., same summarization prompt). Use an in-memory TTL cache.

import time
from collections import defaultdict

_cache: dict[str, tuple[float, str]] = {}
CACHE_TTL = 60.0

def cached_complete(task_text: str) -> str:
    key = task_text.strip().lower()
    if key in _cache:
        ts, val = _cache[key]
        if time.time() - ts < CACHE_TTL:
            return val
    out = complete(task_text)
    _cache[key] = (time.time(), out)
    return out

This avoids redundant classifier calls and model spend for identical inputs within the TTL window.

5. Run the router end-to-end

Wire a small batch through the router best model per task and print which model served each.

requests = [
    "Summarize: The quick brown fox jumps over the lazy dog.",
    "Solve: integrate x^2 from 0 to 3",
    "Refactor this Python loop into a list comprehension: [x for x in range(10) if x%2]",
    "Tell me a joke about compilers",
]

for req in requests:
    t = classify(req)
    route = ROUTES[t]
    print(f"[{t.value}] -> primary {route.primary}")
    print(cached_complete(req)[:80], "...\n")

Expected console output (model names may vary by availability):

[summarize] -> primary google/gemini-flash-1.5
The text describes a quick brown fox jumping over a lazy dog. ...

[math] -> primary anthropic/claude-3-5-sonnet
The integral of x^2 from 0 to 3 is 9. ...

[code] -> primary meta-llama/codellama-34b-instruct
You already have a list comprehension; a loop version would be: ...

[chat] -> primary openai/gpt-4o-mini
Why did the compiler break up? It couldn't resolve its dependencies. ...

6. Track per-token usage metering

A router best model per task is useless without cost visibility. Capture usage from each response and emit a structured log.

import logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("router")

def complete_with_metering(task_text: str) -> str:
    t = classify(task_text)
    route = ROUTES[t]
    resp = client.chat.completions.create(
        model=route.primary,
        messages=[{"role": "user", "content": task_text}],
        max_tokens=route.max_tokens,
    )
    u = resp.usage
    logger.info({
        "task": t.value,
        "model": route.primary,
        "prompt_tokens": u.prompt_tokens,
        "completion_tokens": u.completion_tokens,
    })
    return resp.choices[0].message.content

If your gateway provides per-token usage metering aggregated across providers, you can replace the local log with its usage endpoint for reconciliation.

7. Extend to multi-step agents

For agent loops, classify once per step. Attach the TaskType to the agent state so subsequent calls skip re-classification.

@dataclass
class AgentStep:
    text: str
    task: TaskType

def plan_agent_steps(prompt: str) -> list[AgentStep]:
    # stub: a real planner would decompose; here we classify the whole prompt
    return [AgentStep(prompt, classify(prompt))]

for step in plan_agent_steps("Write a function then explain it"):
    out = complete(step.text, step.task)
    print(step.task, "->", out[:60])

This pattern scales: add TaskType entries for embed, rerank, or vision and extend ROUTES. The router best model per task becomes the backbone of a multi-model agent architecture rather than a one-off shortcut.

Operational notes

• Set timeouts on the client (timeout=10) so a degraded primary fails fast to fallback.
• Rotate fallback models to avoid hotspotting the same cheap model.
• Store ROUTES in a dynamic config store; hot-reload on change.
• The classifier itself is a model call. For high-QPS services, replace it with a local embedding similarity check against task centroids.

Building the router this way gives you explicit control over model selection per task type, measurable spend, and a clean seam to add new models as they ship.