Build a Routed RAG Chatbot with LlamaIndex and Chainlit

This guide is adapted from NVIDIA GenerativeAIExamples app.py, published under the Apache-2.0 licence. Code blocks are reproduced exactly from the source; all prose is original.

---

A naive RAG system routes every query through a single retrieval pipeline regardless of whether the question concerns internal documentation, a product catalogue, a live API, or general world knowledge. That works well enough for demos but collapses under real workloads: precision drops, latency spikes because every query pays the full retrieval tax, and the codebase becomes a monolithic blob that nobody wants to touch when a new knowledge domain is added. Routed RAG solves this by placing a classification layer in front of independent retrieval pipelines and dispatching each query to the one source most likely to answer it correctly.

Engineers who maintain multi-product platforms, enterprise knowledge bases, or domain-specific assistants — anywhere two or more corpora live side by side — will find this pattern immediately useful. Each source gets its own embedding index, retrieval parameters, and prompt template, so you can tune or replace any one without touching the others. The NVIDIA example here pairs the routing logic with a QueryFlow workflow object built on LlamaIndex and surfaces the whole thing through a Chainlit UI, which handles session management and streaming out of the box.

A single GPU capable of running Mistral Large 2 inference — or an API-backed alternative — is sufficient. The Chainlit frontend is CPU-only. NIM endpoints that have been idle can take up to a minute to warm a model from a cold state; the timeout parameters in the source account for this. If you are evaluating hosted inference rather than self-hosting, API spend will be the dominant cost, not hardware. For context on how GPU-aware routing behaves under production traffic, NVIDIA's work on SageMaker inference gateways is worth reading alongside this guide.

---

Prerequisites

• Python 3.10 or later
• A running NVIDIA NIM endpoint, or an NVIDIA API key for hosted inference
• chainlit, llama-index-core, and python-dotenv installed
• The workflow.py module from the NVIDIA GenerativeAIExamples repo containing the QueryFlow class
• Environment variables set in a .env file: at minimum the NIM endpoint URL and API credentials

---

Step 1: Bootstrap the application and configure the top-level workflow

The file opens by importing every dependency and immediately constructing a module-level QueryFlow instance. This object is diagnostic: it exists before any user session starts, so the server process can fail fast during startup if something is misconfigured rather than waiting until a user's first message arrives.

# SPDX-FileCopyrightText: Copyright (c) 2023-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
import time

import chainlit as cl
from dotenv import load_dotenv
from llama_index.core import Settings

from workflow import QueryFlow

load_dotenv()

workflow = QueryFlow(timeout=45, verbose=False)

timeout=45 is intentionally shorter than the per-session value set later — this instance never serves real traffic. load_dotenv() must be called before any LlamaIndex settings are resolved, because Settings reads API keys from environment variables at import time.

---

Step 2: Initialise a fresh workflow per user session

Chainlit calls on_chat_start once per browser session. The critical decision here is that each session gets its own QueryFlow instance rather than sharing the module-level one. Independent instances mean independent state: chat history, any in-memory index caches, and in-flight async tasks cannot bleed between users.

@cl.on_chat_start
async def on_chat_start():

    cl.user_session.set("message_history", [])

    workflow = QueryFlow(timeout=90, verbose=False)

    cl.user_session.set("workflow", workflow)

timeout=90 versus 45 at module level is not cosmetic. NIM endpoints that have been idle can enter a cold state; a timeout=45 guard in a user-facing handler will surface an asyncio.TimeoutError before the model finishes loading, producing a broken streaming response in the UI. Storing the workflow in cl.user_session rather than a global variable is what keeps concurrent sessions independent.

---

Step 3: Define guided conversation starters

The set_starters hook populates the UI with suggested prompts before the user types anything. Each cl.Starter is a labelled shortcut tied to a specific message string. The four prompts are chosen to exercise all four routing branches in QueryFlow: general generation, code generation, NIM-specific documentation retrieval, and BioNemo use-case retrieval. They double as a lightweight smoke test you can run manually after every deployment.

@cl.set_starters
async def set_starters():
    return [
        cl.Starter(
            label="Write a haiku about CPUs",
            message="Write a haiku about CPUs.",
            icon="/avatars/servers",
        ),
        cl.Starter(
            label="Write Docker Compose",
            message="Write a Docker Compose file for deploying a web app with a Redis cache and Postgres database",
            icon="/avatars/screen",
        ),
        cl.Starter(
            label="What NIMs are available?",
            message="Summarize the different large language models that have NVIDIA inference microservices (NIMs) available for them. List as many as you can.",
            icon="/avatars/container",
        ),
        cl.Starter(
            label="Summarize BioNemo use cases",
            message="Write a table summarizing how customers are using bionemo. Use one sentence per customer and include columns for customer, industry, and use case. Make the table between 5 to 10 rows and relatively narrow.",
            icon="/avatars/dna",
        ),
    ]


@cl.on_chat_end
def end():
    logging.info("Chat ended.")

---

Step 4: Handle incoming messages with streaming and source tracking

The main handler is where routing pays off from the user's perspective. It passes the raw query plus accumulated chat history to workflow.run(), which internally classifies the query, selects the appropriate retrieval pipeline, fetches relevant chunks, and returns both a streaming async generator and the source nodes that were used. The streaming loop increments a token counter on every chunk, giving you a cheap throughput metric without a separate observability layer.

@cl.on_message
async def main(user_message: cl.Message, count_tokens: bool = True):
    """
    Executes when a user sends a message. We send the message off to the LlamaIndex chat engine
    for a streaming answer. When the answer is done streaming, we go back over the response
    to identify the sources used, and then add a block of text about the sources.
    """

    msg_start_time = time.time()
    logging.info(f"Received message: <{user_message.content[0:50]}...> ")
    message_history = cl.user_session.get("message_history", [])

    # In case the chat workflow needs extra time to start up,
    # we block until it's ready.

    assistant_message = cl.Message(content="")

    token_count = 0
    with cl.Step(name="Mistral Large 2", type="tool"):

        response, source_nodes = await workflow.run(
            query=user_message.content,
            chat_messages=message_history,
        )

        async for chunk in response:
            token_count += 1
            chars = chunk.delta
            await assistant_message.stream_token(chars)
            
        msg_time = time.time() - msg_start_time
        logging.info(f"Message generated in {msg_time:.1f} seconds.")

    message_history += [
        {"role": "user", "content": user_message.content},
        {"role": "assistant", "content": assistant_message.content},
    ]

    cl.user_session.set("message_history", message_history)

    await assistant_message.send()

The cl.Step context manager wraps the entire retrieval-and-generation cycle as a named step in the Chainlit UI. Naming it after the model makes it immediately obvious in the trace view which model answered which query — useful when debugging a routing mistake. await assistant_message.send() must come after the streaming loop completes; calling it inside the loop would emit partial messages as independent bubbles.

---

Routing options at a glance

---

What to watch out for

Session state and the module-level workflow. The module-level QueryFlow instance (timeout=45) is never used for serving messages — each handler fetches the session-scoped one. If you accidentally pass the module-level instance to main, you will get shared mutable state across all concurrent users, producing subtle history corruption that only appears under load.

Routing misclassification is silent by default. When QueryFlow routes a query to the wrong pipeline, it does not raise an exception. The response still streams — it is just grounded in the wrong corpus. The source_nodes return value is your best diagnostic: log which index was queried for every message during the first week of operation. Misroutes to a general corpus produce fluent but confidently wrong answers, which are harder for users to detect than obvious failures.

Timeout tuning is not optional. NIM endpoints that have not received traffic in several minutes can enter a cold state. A timeout=45 guard in a user-facing handler will surface an asyncio.TimeoutError before the model has finished loading, producing a broken streaming response in the UI. If you are on shared infrastructure or using autoscaling endpoints, lean toward timeout=120 and add explicit retry logic in QueryFlow.run(). For a contrasting approach to cold-start behaviour in managed inference, AWS AgentCore's snapshot memory work is worth reading.

Chat history grows unbounded. The message_history list appended in main has no cap. Over a long session this increases the prompt size on every turn, raising both latency and cost. For production, add a sliding window or summarisation step before passing chat_messages to workflow.run().

count_tokens is not actually wired. The count_tokens: bool = True parameter in main is declared but the token count is never branched on — only logged implicitly. If you want token-gated behaviour such as rejecting queries over a budget, you will need to add that logic yourself.

Starter prompts expose your routing boundaries. The four starters in set_starters are a public map of your knowledge domains. Adversarial users can use them to probe which corpora exist and craft queries designed to extract information from a corpus they should not reach. Treat starter prompt design as part of your security surface, not just UX polish.

---

Where to go next

workflow.py is where the routing logic actually lives — that is the next file to read and modify. Adding a new knowledge domain means registering a new retrieval step in QueryFlow and adding a corresponding starter prompt. For multi-agent orchestration patterns that extend naturally from this routing foundation, DoorDash's multi-agent feature-flag work shows how routing decisions compose with downstream agent calls. For evaluation of retrieval quality per route, integrate a structured grader against held-out question sets for each domain before promoting a new index to production — the approach described in Claude Code's six-grader CI evaluation system maps directly to this use case.