Machine Learning

Wiring ONNX Runtime GenAI into a LangChain RAG Pipeline

By Omer DemirJun 13, 2025⏱ 5 min read👁 — reads⧉ — copies

In part one we picked ONNX Runtime GenAI as the inference engine for an offline, Windows-based RAG system. This post covers the integration: exposing the engine to LangChain and building the optimized model.

Why a custom LLM wrapper

LangChain’s RAG chains expect an LLM that implements its BaseLLM (or BaseChatModel) interface. ONNX Runtime GenAI has its own Generator API, so we bridge the two with a thin adapter that owns the tokenizer and model and implements _generate.

def _generate(self, prompts, stop=None, run_manager=None, **kwargs):
    from onnxruntime_genai import GeneratorParams, Generator

    text_generations: list[str] = []
    answer: str = ""

    # Encode the batch of prompts with the model's tokenizer.
    input_token = self.tokenizer.encode_batch(prompts)

    model_params = self._default_params
    model_params.update(self.model_kwargs)

    # Build generator params (temperature, max_length, top_p, etc.).
    params = GeneratorParams(self.model)
    params.set_search_options(**model_params)
    generator = Generator(self.model, params)

    generator.append_tokens(input_token)
    while not generator.is_done():
        generator.generate_next_token()
        new_token = generator.get_next_tokens()[0]
        answer += self.tokenizer.decode(new_token)

        # Stream to stdout as tokens arrive.
        print(self.tokenizer_stream.decode(new_token), end="", flush=True)

    text_generations.append(answer)
    del generator

    return LLMResult(
        generations=[[Generation(text=text) for text in text_generations]]
    )

A few things worth calling out:

• Token streaming falls out naturally from the while not generator.is_done() loop. You get every token as it’s produced, which lets you stream to a UI or log latency per token.
• del generator matters. Freeing the generator between calls keeps GPU memory flat across a long-running server — important when you’re pinned to a single on-prem box.
• Search options (max_length, temperature, top_p) come straight through set_search_options, so the same wrapper serves greedy decoding for extraction and sampled decoding for summarization.

Once this wrapper exists, the rest of the RAG chain — retriever, prompt template, guardrails — is ordinary LangChain. The engine is invisible to the chain.

The class around _generate

_generate is the hot path, but LangChain needs a bit of scaffolding around it — an _llm_type, default search options, and the model/tokenizer loaded once at construction:

from onnxruntime_genai import Model, Tokenizer
from langchain_core.language_models.llms import BaseLLM

class OnnxRuntimeGenAILLM(BaseLLM):
    model_path: str
    model_kwargs: dict = {}

    def __init__(self, model_path: str, **kwargs):
        super().__init__(model_path=model_path, **kwargs)
        self.model = Model(model_path)              # loaded once, reused per request
        self.tokenizer = Tokenizer(self.model)
        self.tokenizer_stream = self.tokenizer.create_stream()

    @property
    def _llm_type(self) -> str:
        return "onnxruntime-genai"

    @property
    def _default_params(self) -> dict:
        # Greedy by default; override per call for sampled decoding.
        return {"max_length": 512, "temperature": 0.0, "top_p": 1.0}

Loading Model and Tokenizer once at startup — not per request — is what keeps first-token latency flat. Everything a request needs is already resident on the GPU.

Building the optimized ONNX model

The throughput numbers from part one depend on running an optimized graph, not a naive export. ONNX Runtime GenAI ships a model builder that fuses operators and applies the optimizations Phi-3 needs:

onnxruntime-genai.builder --model phi-3-mini --output ./phi3_optimized.onnx

We run this as a one-time conversion job and ship the resulting artifact with the application. Because the target environment is air-gapped, baking the optimized model into the deployment package is not just convenient — it’s required. There’s no model download at runtime.

The builder emits more than a single .onnx file — you get the quantized weights, the tokenizer, and a genai_config.json describing the graph and default search options:

phi3_optimized/
├── model.onnx              # optimized, fused graph
├── model.onnx.data         # INT4 block-quantized weights
├── genai_config.json       # EP config + default search options
└── tokenizer.json

The two knobs that matter most are precision (int4 for the bandwidth win) and execution provider (cuda on the A100). Those two choices are exactly what produced part one’s throughput and latency numbers.

Assembling the RAG chain

With the LLM wrapped, the chain is unremarkable LangChain — which is the point:

llm = OnnxRuntimeGenAILLM(model_path="./phi3_optimized")
retriever = local_vector_store.as_retriever(search_kwargs={"k": 4})

prompt = ChatPromptTemplate.from_template(
    "Answer using only the context.\n\nContext:\n{context}\n\nQuestion: {question}"
)

rag_chain = (
    {"context": retriever | format_docs, "question": RunnablePassthrough()}
    | prompt
    | llm            # the ONNX Runtime GenAI wrapper
    | guardrails     # validate before returning
)

Putting it together

The end-to-end flow on the Windows box looks like this:

1. Load the optimized Phi-3 ONNX model and tokenizer once at startup.
2. Retrieve context from the local vector store for the user’s query.
3. Compose the prompt with retrieved context via a LangChain prompt template.
4. Generate through the custom BaseLLM, streaming tokens as they arrive.
5. Guard the output with guardrails before returning it.

With the optimized model and the wrapper above, the full pipeline stayed under the 5-second budget even with retrieval-inflated prompts — the goal we set at the start of the series.

Here’s roughly where that budget goes on a typical query (retrieved context ~256 tokens, 256-token answer):

Generation is ~80% of the wall-clock, which is why the engine choice from part one mattered so much: shave 20% off decode throughput and the whole budget tightens.

Takeaways

• A ~40-line adapter is all it takes to make ONNX Runtime GenAI look like any other LangChain LLM.
• Always build the optimized graph with onnxruntime-genai.builder; the raw export leaves most of the performance on the table.
• For air-gapped deployments, treat the optimized model as a build artifact and ship it with the app.

References

• ONNX Runtime GenAI documentation
• LangChain custom LLM guide

#slm#onnx#rag#langchain#windows