Multi-model RAG: routing retrieval vs generation calls

Most RAG stacks bolt a single chat model onto a vector search and call it a day. Multi-model RAG routing treats retrieval and generation as separate inference problems, letting you pick a cheap embedder or reranker for search and a stronger model for synthesis. This guide lays out an ordered path to split those calls without adding operational drag.

1. Separate retrieval and generation concerns

Retrieval is a ranking problem. Generation is a conditional text problem. Binding both to one model forces compromises: you either overpay for embedding-quality LLM calls or underpower your answers with a weak chat model.

The first move is to draw a hard line in your code. Embeddings and reranking happen before any user-facing token is generated. The generation step consumes the filtered context and produces the answer.

from openai import OpenAI

client = OpenAI(base_url="https://api.openai.com/v1", api_key="sk-...")

# Retrieval: embed query with a dedicated embedding model
q_emb = client.embeddings.create(
    model="text-embedding-3-small",
    input="How do I rotate API keys without downtime?"
).data[0].embedding
# Vector DB search returns candidate chunks (omitted)

Pitfall: teams often call the generation LLM to “score” chunks. At 500 chunks per query, that’s 500 completions. Use an embedding model or a local cross-encoder instead.

2. Choose a retrieval model independent of generation

Pick the smallest embedding model that holds recall for your domain. OpenAI’s text-embedding-3-small is 1536-dim and cheap; text-embedding-3-large adds precision at higher cost. Run offline recall eval before committing.

For noisy corpora, add a reranker. A small instruction-tuned model can prune top-20 to top-5 with one token of output per candidate.

Rerank with a small model

def rerank(query, chunks, client):
    scored = []
    for c in chunks:
        resp = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": "Score relevance 0-9."},
                {"role": "user", "content": f"Q: {query}\nDoc: {c[:500]}"}
            ],
            max_tokens=1,
            temperature=0
        )
        try:
            score = int(resp.choices[0].message.content.strip() or 0)
        except ValueError:
            score = 0
        scored.append((score, c))
    return [c for _, c in sorted(scored, reverse=True)[:5]]

Tradeoff: synchronous reranking adds latency linear to candidate count. Batch the calls with asyncio or use a dedicated rerank endpoint if your gateway exposes one.

3. Route generation by query class

Multi-model RAG routing pays off when you stop sending every query to your most expensive model. Classify the query intent first.

A heuristic works surprisingly well:

def route_generation(query):
    words = query.split()
    if len(words) > 12 or "compare" in query.lower() or "why" in query.lower():
        return "gpt-4o"
    return "gpt-4o-mini"

Then execute against the routed model:

model = route_generation(user_q)
answer = client.chat.completions.create(
    model=model,
    messages=[
        {"role": "system", "content": "Answer strictly from context."},
        {"role": "user", "content": f"{context}\n\n{user_q}"}
    ]
)

For production, replace the heuristic with a 10-line classifier fine-tune or a prompt to a tiny model. The point is that the routing decision is explicit and logged.

Common mistake: routing on latency alone. A cheap model that returns “I don’t know” triggers a retry to the expensive model, doubling cost. Route on predicted complexity, not just tail latency.

4. Implement fallback without breaking context

Providers throttle. When gpt-4o is rate-limited, a silent failure is worse than a degraded answer. Use a gateway that honors client routing directives and provides automatic fallback when a provider is degraded.

An OpenAI-compatible gateway like n4n.ai lets you target one endpoint covering 240+ models and declare fallbacks per call, so the same request shape degrades to a backup model without code branches.

resp = client.chat.completions.create(
    model="anthropic/claude-3-5-sonnet",
    messages=msgs,
    extra_headers={
        "x-fallback-models": "openai/gpt-4o,meta/llama-3-70b-instruct"
    }
)

Pitfall: fallback models have different prompt tolerances. Claude expects system prompts in a specific field; Llama variants may need different instruction formatting. Keep your system prompt neutral and test the fallback path quarterly.

5. Forward cache-control and meter per stage

Retrieval embeddings for repeated queries are identical. Generation prompts with static corpus context can use provider prompt caching. Forward cache hints explicitly:

client.chat.completions.create(
    model="gpt-4o",
    messages=msgs,
    extra_headers={"cache-control": "ephemeral"}
)

Without forwarding, a gateway may strip the hint and you pay full input tokens every call.

Meter usage separately for retrieval and generation. The usage object gives per-call tokens:

print(resp.usage.prompt_tokens, resp.usage.completion_tokens)

If you aggregate across providers, tag each call with a stage field in your logging. Multi-model RAG routing only saves money if you can see that reranking costs 0.2% of total tokens while generation eats 80%.

6. Observe and tune the routing table

Set up a minimal dashboard: p95 latency per model, token cost per stage, and answer acceptance rate (thumbs up/down or heuristic length check). Review weekly.

Start with conservative routing: route 90% of traffic to the cheap model, 10% to strong. Measure quality delta. If acceptance is within 2%, expand cheap routing.

Routing table example

{
  "retrieval_embed": "text-embedding-3-small",
  "rerank": "gpt-4o-mini",
  "generation_default": "gpt-4o-mini",
  "generation_complex": "gpt-4o",
  "fallback_chain": ["claude-3-5-sonnet", "llama-3-70b-instruct"]
}

Treat this as code. Version it. Roll back if quality drops.

Common pitfalls

• Coupled version upgrades. Bumping the generation model without re-evaluating embedding recall silently degrades answers.
• Synchronous rerank loops. Calling a model per chunk blocks the request thread. Use async or batch.
• Fallback prompt drift. Backup models receive the same messages but parse them differently. Test the degraded path.
• No per-stage metering. You can’t tune what you don’t measure. Log retrieval vs generation tokens distinctly.
• Over-routing to cheap models. Saving $0.0002 per call isn’t worth a 5% drop in task success.

Multi-model RAG routing is not a one-time config. It’s a continuous practice of matching model capability to subtask and watching the meters. Do that, and your pipeline stays both cheap and correct under load.