Skip to main content

Model Routing for RAG Systems: A Practical Guide

Vishal Vishwakarma10 min read

Retrieval-augmented generation (RAG) settled quickly into a standard architecture: embed and index your documents, retrieve what is relevant to a query, assemble context, and hand the result to a language model to generate the answer. Most of the engineering attention goes to the retrieval half, chunking strategies, embedding choices, rerankers, and that attention is well spent. But the line item on the bill, and the quality ceiling on the answer, both live in the half that gets chosen once and forgotten: which model generates.

The overlooked property of RAG generation traffic is its spread. The same deployment serves "what is the refund window" and "do these three clauses, retrieved from different policy versions, contradict each other," and those two queries need very different machinery. Answering the first with a frontier model wastes most of what you paid for; answering the second with a small model produces a confidently wrong synthesis. This guide covers the practical integration of a routing layer at exactly that step, using Inferbase as the working example: what changes in your pipeline (almost nothing), what decisions you own (the pool and the privacy posture), and how to verify the routing is doing what it claims.

Where the boundary sits

The whole integration reduces to one statement: retrieval stays yours, generation gets routed. Your embeddings, your vector database, your chunking and ranking logic are untouched, and they never interact with the routing layer. The only call that moves is the one your pipeline already makes last, the assembled prompt going to a model.

A RAG pipeline diagram showing where routing sits. Inside a dashed container labeled your stack, unchanged: a user query of any difficulty flows to a retriever using your embeddings and index, which queries your vector database, and the results flow into a context assembly step that combines retrieved chunks and the question into one prompt, with a note that retrieved PII transits here, which is why the generation step owns the privacy decision. One highlighted call, model set to auto, crosses to the Inferbase generation layer, where a simple refund-window question routes to a small fast model and a multi-document clause-reconciliation question routes to a frontier model, with the pool restricted to models your review approved and every response reporting which model served and why. The answer returns to your stack. A footnote states that in classic RAG, retrieval stays entirely in your stack and routing applies to the generation call, where per-query difficulty actually varies.

In code, the change is confined to the client construction and the model parameter of your existing generation call. A representative pipeline tail:

from openai import OpenAI

client = OpenAI(
    api_key="inf_your_api_key",
    base_url="https://api.inferbase.ai/api/v1/inference",
)

def answer(query: str, chunks: list[str]) -> str:
    context = "\n\n---\n\n".join(chunks)
    response = client.chat.completions.create(
        model="auto",
        messages=[
            {"role": "system", "content": (
                "Answer strictly from the provided context. "
                "If the context does not contain the answer, say so."
            )},
            {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {query}"},
        ],
        extra_body={
            "routing": {
                "optimize": "balanced",
                "min_context_window": 32000,
            }
        },
    )
    return response.choices[0].message.content

Everything upstream of answer() is your existing pipeline. The min_context_window constraint is the one routing parameter RAG almost always wants set, because your assembled context has a size floor and the router should never pick a model that cannot hold it; size it to your retrieval configuration (chunk size times top-k, plus headroom for the question and the answer).

Why the spread is worth money

Consider a support knowledge-base deployment at moderate scale, one million queries a month, averaging 3,000 prompt tokens (retrieved context dominates) and 300 completion tokens. Traffic analysis in deployments like this consistently shows a shape rather than a uniform mass: most queries are single-fact lookups the retrieval step has already essentially answered, and a minority require genuine cross-document reasoning.

Traffic sliceShareWhat it needsIndicative price tier
Single-passage lookup ("what is X")~60-70%Small fast model$0.10-0.50 per 1M tokens
Standard synthesis (one topic, few chunks)~20-30%Mid-tier model$0.50-3 per 1M tokens
Cross-document reasoning, edge cases~10-15%Frontier model$3-15 per 1M tokens

Sending everything to the frontier tier prices the entire million queries at the hardest slice's rate. Routing prices each slice at its own tier, and because the cheap slice is the large one, the arithmetic moves materially: the blended per-query cost lands near the cheap end rather than the expensive end, typically a several-fold reduction on the generation line item at this traffic shape. The exact figure depends on your spread, which is the point: the router's job is to discover your spread per query instead of you guessing it once for all queries. (Tier prices move constantly; the model catalog has current numbers, and the structure of the argument survives any particular price list.)

The honest counterpart: if your RAG traffic genuinely is uniform, an internal legal-research tool where every query is hard, routing sends most requests to the strong tier and the savings are modest. Routing monetizes variance. It does not manufacture it.

The privacy decision lives at this step

Context assembly is where retrieved content, including whatever personal or regulated data your documents contain, gets placed into a prompt and leaves your stack. That makes the generation call, not retrieval, the step your compliance review actually cares about, and it is worth handling explicitly rather than inheriting whatever the default is.

Two controls map onto it directly. The first is the model pool: a curated list of model ids, set on the API key, that the router may never leave. The approval conversation ("which models may see claims data") happens once, when the pool is agreed, and every subsequent request is mechanically confined to it:

# Set once on the API key for the regulated surface; requests inherit it.
# Per-request override exists but is better reserved for exceptions.
routing = {
    "optimize": "quality",
    "model_pool": [
        "openai/gpt-oss-120b",
        "Qwen/Qwen3-235B-A22B-Instruct-2507-tput",
    ],
}

The second is retention. Inferbase does not store prompt or response content on the serve path in the first place; usage logs and routing audit records are metadata only. The two features that can store content, saved playground conversations and eval uploads, are opt-in, and an account-level zero-retention flag disables both outright. For a RAG surface carrying PII, flip that flag before the first production request, and the retrieved content's exposure reduces to the transient generation call against models your pool already approved.

Wiring it through LangChain and LlamaIndex

Most production RAG stacks call the model through a framework rather than the raw SDK, and both major frameworks accept an OpenAI-compatible endpoint without ceremony. The integration point is identical in spirit: point the chat model at the routed base URL and set the model to auto or a pin.

LangChain's ChatOpenAI takes the base URL and an extra_body for the routing parameters, so the rest of your chain, retriever, prompt template, output parser, is untouched:

from langchain_openai import ChatOpenAI

llm = ChatOpenAI(
    model="auto",
    api_key="inf_your_api_key",
    base_url="https://api.inferbase.ai/api/v1/inference",
    extra_body={
        "routing": {
            "optimize": "balanced",
            "min_context_window": 32000,
        }
    },
)
# rag_chain = ({"context": retriever | format_docs, "question": ...} | prompt | llm | parser)

LlamaIndex routes OpenAI-compatible endpoints through its OpenAILike class, which skips the model-name validation its native OpenAI class applies. Mark it as a chat model and pass routing parameters the same way, inside extra_body:

from llama_index.llms.openai_like import OpenAILike

llm = OpenAILike(
    model="auto",
    api_base="https://api.inferbase.ai/api/v1/inference",
    api_key="inf_your_api_key",
    is_chat_model=True,
    additional_kwargs={
        "extra_body": {"routing": {"optimize": "balanced", "min_context_window": 32000}}
    },
)

One boundary note that applies to both frameworks: this wires the generation model. Your embedding model, wherever it runs today, stays where it is, since embeddings are part of the retrieval half that never crosses the routing boundary.

Sizing min_context_window from your retrieval configuration

The min_context_window constraint deserves actual arithmetic rather than a guess, because both failure directions cost you. Set it too low and the router may pick a model your assembled context does not fit, failing the request; set it needlessly high and you exclude the cheap models that were the point of routing.

The inputs are all in your retrieval configuration. A worked example with common numbers:

ComponentSizingTokens
Retrieved context8 chunks x 512 tokens4,096
Chunk separators and framing~10 per chunk80
System promptfixed250
User questionp95, not average300
Response headroom (max_tokens)your cap1,500
Safety margin~15% of the above900
Total floor~7,100

That deployment needs min_context_window of about 8,000, not the 32,000 in the earlier snippets, and the difference matters: an 8K floor keeps small fast models eligible for the easy majority of queries, while a reflexive 32K floor would have excluded some of them for no operational reason. Recompute the floor when you change chunk size, top-k, or the response cap, since those three move it most. If you later add reranking that trims to fewer, better chunks, lower the floor accordingly and the router's cheap tier gets correspondingly wider.

Verifying the routing, not trusting it

A router that cannot show its work is a liability in a pipeline that answers from your documents. Every routed response disclosed which model served it: non-streaming responses carry it in the model field and the X-Inferbase-Model header; streaming responses open with a routing event before the first token, including the task and complexity classification, the served model, scored alternatives, and a decisiveness figure for how clear-cut the pick was.

The verification worth doing in the first week is joining that disclosure against your own query logs. Sample the queries classified simple and confirm they are in fact lookups; sample the complex ones and confirm they are the multi-document cases. In our experience the disagreements you find are as often retrieval problems as routing problems, a "simple" classification on a query your retriever failed, handing the model thin context, reads as a routing miss but is actually a recall miss upstream. The routing metadata gives you a second, independent signal on pipeline health, which is a quiet benefit beyond the cost arithmetic.

Two failure-handling notes complete the picture. Model and provider failures inside the routed call are handled by the same fallback mechanics that apply to any served request, under one request deadline, with the substitution disclosed. And if a query class must never be routed at all, pin it: model accepts a concrete id per request, and a mixed deployment (auto for the general lane, pinned for contractual flows) is a normal end state, not a compromise.

How to add routing to an existing RAG deployment, step by step

  1. Measure your spread first. Classify a week of production queries by hand or with a cheap model into lookup / synthesis / reasoning buckets. If the distribution is 90% one bucket, routing will help less than this guide implies, and you should know that before integrating anything.
  2. Curate the pool with whoever owns data approval, sized to your context floor (min_context_window) and your privacy constraints. Set zero-retention if the documents warrant it.
  3. Swap the generation client (two lines) on a single surface, optimize: "balanced", and log the routing event per request.
  4. Run the week-one verification above: served-model distribution against query difficulty, decisiveness trend, per-request cost from the usage events.
  5. Then tune: move the surface to cost or quality if the balanced distribution suggests headroom, add pins for the flows that need them, and only then extend to the next surface.

The API reference documents the full routing parameter surface. The retrieval half of your pipeline, meanwhile, has not noticed anything changed, which is the property that makes this integration worth trialing in an afternoon rather than scheduling for a quarter.

Frequently asked questions

RAGmodel routingretrieval-augmented generationinferenceintegration

Have thoughts on this article?

We would love to hear your feedback, questions, or experience with these topics. Reach out on social media or drop us a message.

Related Articles

Stay up to date

Get notified when we publish new articles on AI model selection, cost optimization, and infrastructure planning.

Your AI stack shouldn't stand still.

Every month new models become cheaper, faster, and more capable. Inferbase ensures your application automatically benefits without changing a single API call.