Stream RAG responses in LlamaIndex with n4n.ai

Streaming RAG responses in LlamaIndex lets users see tokens as they arrive instead of waiting for the full generation. This llamaindex streaming rag responses tutorial walks through wiring a query engine to n4n.ai’s OpenAI-compatible endpoint so you get token-by-token output with automatic provider fallback and per-token usage metering built in.

Step 1: Install dependencies

Start with a clean virtual environment and install the minimal set of packages. You need LlamaIndex core, the OpenAI integration (since n4n.ai speaks the OpenAI API), and a vector store — here we use Chroma for simplicity.

python -m venv .venv
source .venv/bin/activate
pip install -U llama-index llama-index-llms-openai llama-index-vector-stores-chroma chromadb

Verify the install:

python -c "import llama_index; print(llama_index.__version__)"

You should see a version string like 0.11.x or newer.

Step 2: Configure the n4n.ai client

n4n.ai exposes a single OpenAI-compatible base URL. Point the LlamaIndex OpenAI wrapper at it and pass your n4n.ai API key. The gateway handles model routing, fallback, and usage accounting transparently.

# config.py
import os
from llama_index.llms.openai import OpenAI

N4N_API_KEY = os.getenv("N4N_API_KEY")
if not N4N_API_KEY:
    raise RuntimeError("Set N4N_API_KEY in your environment")

llm = OpenAI(
    model="gpt-4o-mini",           # logical model name; n4n.ai maps to an available provider
    api_key=N4N_API_KEY,
    api_base="https://api.n4n.ai/v1",
    temperature=0.1,
    max_tokens=1024,
    streaming=True,                # critical: enables token streaming
)

Set the key in your shell before running anything:

export N4N_API_KEY="n4n_sk_..."

Step 3: Ingest documents into a vector index

Create a small corpus, embed it, and persist to Chroma. In production you’d swap the embedding model and chunking strategy, but this gets you a working index in under 30 lines.

# ingest.py
import os
from pathlib import Path
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, StorageContext
from llama_index.vector_stores.chroma import ChromaVectorStore
import chromadb
from config import llm

# 1. Load documents
documents = SimpleDirectoryReader("./data").load_data()
print(f"Loaded {len(documents)} documents")

# 2. Initialize Chroma
chroma_client = chromadb.PersistentClient(path="./chroma_db")
chroma_collection = chroma_client.get_or_create_collection("rag_demo")
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
storage_context = StorageContext.from_defaults(vector_store=vector_store)

# 3. Build and persist index
index = VectorStoreIndex.from_documents(
    documents,
    storage_context=storage_context,
    show_progress=True,
)
print("Index built and persisted to ./chroma_db")

Create a data/ directory with a few .txt or .md files, then run:

python ingest.py

Success indicator: the script prints the document count and “Index built and persisted” without errors. You can inspect the Chroma store with chromadb CLI if needed.

Step 4: Build a streaming query engine

LlamaIndex’s QueryEngine supports streaming via the streaming=True flag. The key is passing the already-configured streaming LLM and ensuring the response synthesizer doesn’t buffer.

# query.py
from llama_index.core import VectorStoreIndex, StorageContext, get_response_synthesizer
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.vector_stores.chroma import ChromaVectorStore
import chromadb
from config import llm

# 1. Reconnect to the persisted index
chroma_client = chromadb.PersistentClient(path="./chroma_db")
chroma_collection = chroma_client.get_collection("rag_demo")
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = VectorStoreIndex.from_vector_store(vector_store, storage_context=storage_context)

# 2. Configure a streaming response synthesizer
#    `streaming=True` on the LLM plus `response_mode="tree_summarize"` or "compact"
#    gives token-by-token output without waiting for full retrieval + synthesis.
response_synthesizer = get_response_synthesizer(
    llm=llm,
    response_mode="compact",      # streams each chunk's answer then merges
    streaming=True,
)

# 3. Assemble the query engine
query_engine = RetrieverQueryEngine(
    retriever=index.as_retriever(similarity_top_k=4),
    response_synthesizer=response_synthesizer,
)

# 4. Stream a query
query = "What are the key design principles mentioned in the docs?"
print(f"Query: {query}\n")
print("Response (streaming):\n")

streaming_response = query_engine.query(query)
for token in streaming_response.response_gen:
    print(token, end="", flush=True)

print("\n\n--- Done ---")

Run it:

python query.py

Verification: You should see tokens appear one by one (or in small chunks) in the terminal, not a single block after a long pause. The total latency to first token should be noticeably lower than a non-streaming run.

Step 5: Add a simple CLI for interactive use

A one-shot script is fine for testing, but a REPL makes iteration faster. Wrap the query engine in a loop that respects Ctrl+C and prints usage metadata from the response object.

# cli.py
import sys
from query import query_engine

def main():
    print("LlamaIndex RAG streaming CLI — type 'exit' to quit\n")
    while True:
        try:
            query = input("❯ ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nBye")
            break

        if query.lower() in {"exit", "quit"}:
            break
        if not query:
            continue

        print()
        response = query_engine.query(query)
        for token in response.response_gen:
            print(token, end="", flush=True)
        print("\n")

        # Metadata n4n.ai returns via OpenAI-compatible headers
        if hasattr(response, "metadata") and response.metadata:
            usage = response.metadata.get("usage")
            if usage:
                print(f"  [usage] prompt={usage.get('prompt_tokens')} "
                      f"completion={usage.get('completion_tokens')} "
                      f"total={usage.get('total_tokens')}")

if __name__ == "__main__":
    main()

Run python cli.py and ask a few questions. Verify:

• Tokens stream immediately
• Usage metadata prints after each response (prompt/completion/total tokens)
• No tracebacks on empty input or interrupt

Step 6: Handle provider fallback and rate limits gracefully

One reason to put n4n.ai in front of your LlamaIndex pipeline is automatic fallback when a provider hits rate limits or degrades. The gateway returns standard OpenAI error codes; you just need to catch them and optionally retry with backoff.

# resilient_query.py
import time
from openai import RateLimitError, APIConnectionError
from query import query_engine

def query_with_retry(query_str: str, max_retries: int = 3, base_delay: float = 1.0):
    for attempt in range(max_retries):
        try:
            return query_engine.query(query_str)
        except (RateLimitError, APIConnectionError) as e:
            if attempt == max_retries - 1:
                raise
            delay = base_delay * (2 ** attempt)
            print(f"  [retry {attempt + 1}/{max_retries}] {e.__class__.__name__}: waiting {delay}s")
            time.sleep(delay)

if __name__ == "__main__":
    resp = query_with_retry("Summarize the architecture section")
    for token in resp.response_gen:
        print(token, end="", flush=True)
    print()

Test this by temporarily exhausting a provider quota (or simulate with a mock). The gateway will route to a healthy provider and the stream resumes without your code knowing the switch happened.

Step 7: Wire into a FastAPI endpoint for production

Most teams expose the streaming query engine over HTTP. FastAPI with StreamingResponse is the standard pattern. The endpoint yields tokens as SSE or plain text chunks.

# server.py
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from query import query_engine
import uvicorn

app = FastAPI(title="LlamaIndex RAG Streaming API")

class QueryRequest(BaseModel):
    query: str
    stream: bool = True

@app.post("/query")
async def query_endpoint(req: QueryRequest):
    if not req.query.strip():
        raise HTTPException(400, "Query cannot be empty")

    response = query_engine.query(req.query)

    if req.stream:
        def token_generator():
            for token in response.response_gen:
                yield token
        return StreamingResponse(token_generator(), media_type="text/plain")
    else:
        return {"response": str(response)}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)

Run the server:

python server.py

Test with curl (streaming):

curl -N -X POST http://localhost:8000/query \
  -H "Content-Type: application/json" \
  -d '{"query": "What is the deployment model?", "stream": true}'

The -N flag disables curl’s output buffering so you see tokens as they arrive. Verify:

• First token arrives within ~500ms (network + retrieval + first provider token)
• Stream completes without truncation
• Non-streaming mode ("stream": false) returns a single JSON response

Step 8: Observe per-token usage in logs

n4n.ai forwards provider usage fields in the final response metadata. For streaming, the final chunk includes cumulative token counts. Capture this for cost tracking or quota enforcement.

# usage_logger.py
import json
from query import query_engine

def log_usage(query: str):
    response = query_engine.query(query)
    full_text = ""
    for token in response.response_gen:
        full_text += token
        print(token, end="", flush=True)
    print()

    # Final metadata contains usage
    meta = getattr(response, "metadata", {})
    usage = meta.get("usage")
    if usage:
        log_entry = {
            "query": query,
            "prompt_tokens": usage.get("prompt_tokens"),
            "completion_tokens": usage.get("completion_tokens"),
            "total_tokens": usage.get("total_tokens"),
            "model": meta.get("model"),
        }
        print(json.dumps(log_entry, indent=2))
    else:
        print("No usage metadata returned")

if __name__ == "__main__":
    log_usage("Explain the caching strategy")

Run it and confirm the JSON log prints after the stream finishes. In production you’d ship this to your observability stack (Datadog, Prometheus, etc.).

Step 9: Tune retrieval for lower latency to first token

Streaming feels fast only if retrieval doesn’t dominate. Two quick wins:

1. Reduce similarity_top_k — fewer chunks means less context to stuff into the prompt, faster first token.
2. Use a smaller embedding model — text-embedding-3-small is ~5x faster than text-embedding-3-large with acceptable recall for many domains.

# In query.py, adjust the retriever:
retriever = index.as_retriever(similarity_top_k=3)  # down from 4

Benchmark with a simple timer:

import time
from query import query_engine

start = time.perf_counter()
resp = query_engine.query("What is the API rate limit?")
first_token_time = None
for i, token in enumerate(resp.response_gen):
    if first_token_time is None:
        first_token_time = time.perf_counter()
    print(token, end="", flush=True)
print(f"\nFirst token: {(first_token_time - start)*1000:.0f}ms")
print(f"Total: {(time.perf_counter() - start)*1000:.0f}ms")

Target: first token under 800ms on a warm index. If you’re higher, profile retrieval vs. provider latency separately.

Step 10: Deploy with Docker for consistency

Package the whole stack — app, Chroma, and config — so staging matches production.

# Dockerfile
FROM python:3.11-slim

WORKDIR /app

# System deps for Chroma
RUN apt-get update && apt-get install -y --no-install-recommends \
    gcc libpq-dev && rm -rf /var/lib/apt/lists/*

COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

COPY . .

ENV PYTHONUNBUFFERED=1
EXPOSE 8000

CMD ["python", "server.py"]

# requirements.txt
llama-index==0.11.*
llama-index-llms-openai==0.2.*
llama-index-vector-stores-chroma==0.2.*
chromadb==0.5.*
fastapi==0.112.*
uvicorn==0.30.*
python-dotenv==1.0.*

Build and run:

docker build -t rag-streaming .
docker run -p 8000:8000 -e N4N_API_KEY="$N4N_API_KEY" -v $(pwd)/chroma_db:/app/chroma_db rag-streaming

The volume mount persists the Chroma database across container restarts. Verify the API responds at http://localhost:8000/query.

Common pitfalls

What’s next

• Swap compact response mode for tree_summarize if you need hierarchical synthesis over many chunks — still streams, but with a different merge strategy.
• Add a reranker (Cohere, Jina, or cross-encoder) between retrieval and synthesis to improve answer quality without increasing top_k.
• Instrument the FastAPI endpoint with OpenTelemetry; n4n.ai’s gateway already emits provider-level latency headers you can correlate.
• For multi-tenant apps, namespace Chroma collections per tenant and pass the tenant ID through the request context.

You now have a production-ready streaming RAG pipeline: ingestion, retrieval, token streaming, fallback, usage accounting, and a deployable service. The only thing that changes at scale is the vector store backend and the routing logic — n4n.ai handles the model plumbing so you don’t have to.