← Back to DevBytes

Building a Knowledge Base Chatbot with vLLM: Complete Guide

Building a Knowledge Base Chatbot with vLLM: Complete Guide

Large language models are powerful, but they hallucinate, forget domain specifics, and cannot access your private documentation. A knowledge base chatbot solves this by combining a retrieval pipeline with a high-throughput inference engine. In this guide, we will build a production-grade Retrieval-Augmented Generation (RAG) chatbot using vLLM as the serving backend, an embedding model for semantic search, and a vector database for storage.

What Is vLLM?

vLLM is an open-source inference engine developed at UC Berkeley that dramatically accelerates LLM serving. Its core innovation is PagedAttention, a memory management technique inspired by operating system virtual memory that reduces KV-cache waste and enables much higher throughput than naive Hugging Face Transformers serving.

For a knowledge base chatbot, vLLM matters because:

Why a Knowledge Base Chatbot?

A knowledge base chatbot grounds model responses in your own documents. Instead of relying on what the model memorized during pretraining, you retrieve relevant chunks from a vector store and inject them into the prompt. This gives you:

Architecture Overview

The system has four components: a document ingestion pipeline, an embedding model, a vector database, and the vLLM inference server. The chat flow is: user asks a question, the system embeds the query, retrieves the top-k matching chunks, constructs a prompt with that context, and vLLM generates the answer.

Prerequisites and Setup

You will need a machine with an NVIDIA GPU (at least 16GB VRAM for a 7B model) and Docker installed. We will use Python 3.10+ and several libraries. Create a working directory and install dependencies:

mkdir kb-chatbot && cd kb-chatbot
python -m venv venv && source venv/bin/activate
pip install vllm faiss-cpu sentence-transformers langchain pypdf fastapi uvicorn

For the embedding model, we will use BAAI/bge-small-en-v1.5, which is lightweight and high quality. For the LLM, we will use mistralai/Mistral-7B-Instruct-v0.2, a strong open model that runs comfortably on a single 24GB GPU.

Step 1: Starting the vLLM Server

vLLM ships with a built-in OpenAI-compatible server. Start it in a separate terminal:

python -m vllm.entrypoints.openai.api_server \
  --model mistralai/Mistral-7B-Instruct-v0.2 \
  --port 8000 \
  --max-model-len 8192 \
  --gpu-memory-utilization 0.9 \
  --enable-prefix-caching

The --enable-prefix-caching flag is especially valuable for RAG: since your system prompt and retrieved context are reused across queries, prefix caching avoids recomputing their KV cache, cutting latency significantly. Once running, you can test it with curl:

curl http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "mistralai/Mistral-7B-Instruct-v0.2",
    "messages": [{"role": "user", "content": "Say hello."}]
  }'

Step 2: Building the Document Ingestion Pipeline

First, create a script that loads PDFs, splits them into chunks, embeds them, and stores them in FAISS. Create ingest.py:

import os
import glob
from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS

def load_documents(data_dir="data"):
    docs = []
    for pdf_path in glob.glob(os.path.join(data_dir, "*.pdf")):
        loader = PyPDFLoader(pdf_path)
        docs.extend(loader.load())
    return docs

def split_documents(docs, chunk_size=512, chunk_overlap=64):
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        separators=["\n\n", "\n", ". ", " ", ""]
    )
    return splitter.split_documents(docs)

def build_vectorstore(chunks, persist_dir="vectorstore"):
    embeddings = HuggingFaceEmbeddings(
        model_name="BAAI/bge-small-en-v1.5",
        model_kwargs={"device": "cuda"}
    )
    vectorstore = FAISS.from_documents(chunks, embeddings)
    vectorstore.save_local(persist_dir)
    print(f"Saved {len(chunks)} chunks to {persist_dir}")

if __name__ == "__main__":
    docs = load_documents("data")
    chunks = split_documents(docs)
    build_vectorstore(chunks)

Place your PDF files in a data/ directory and run python ingest.py. The chunk size of 512 tokens balances retrieval precision against context completeness. Smaller chunks retrieve more precisely but may miss surrounding context; larger chunks provide more context but dilute relevance signals.

Step 3: Building the Retrieval and Generation Pipeline

Now create the main chatbot module that retrieves relevant chunks and calls vLLM. Create chatbot.py:

import json
import requests
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS

VLLM_URL = "http://localhost:8000/v1/chat/completions"
MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.2"

SYSTEM_PROMPT = """You are a helpful knowledge base assistant. Answer the user's
question using ONLY the provided context. If the context does not contain the
answer, say "I don't have enough information to answer that." Do not make up facts.
Always cite which document section your answer comes from."""

def load_vectorstore(persist_dir="vectorstore"):
    embeddings = HuggingFaceEmbeddings(
        model_name="BAAI/bge-small-en-v1.5",
        model_kwargs={"device": "cuda"}
    )
    return FAISS.load_local(persist_dir, embeddings, allow_dangerous_deserialization=True)

def retrieve_context(vectorstore, query, k=4):
    results = vectorstore.similarity_search_with_score(query, k=k)
    context_blocks = []
    for i, (doc, score) in enumerate(results, 1):
        source = doc.metadata.get("source", "unknown")
        page = doc.metadata.get("page", "?")
        context_blocks.append(
            f"[Context {i}] (Source: {source}, Page: {page})\n{doc.page_content}"
        )
    return "\n\n".join(context_blocks)

def generate_answer(query, context):
    user_message = f"Context:\n{context}\n\nQuestion: {query}\n\nAnswer:"
    payload = {
        "model": MODEL_NAME,
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_message}
        ],
        "temperature": 0.1,
        "max_tokens": 512,
        "top_p": 0.9
    }
    response = requests.post(VLLM_URL, json=payload)
    response.raise_for_status()
    return response.json()["choices"][0]["message"]["content"]

def chat():
    vectorstore = load_vectorstore()
    print("Knowledge Base Chatbot ready. Type 'quit' to exit.\n")
    while True:
        query = input("You: ").strip()
        if query.lower() in ("quit", "exit"):
            break
        context = retrieve_context(vectorstore, query, k=4)
        answer = generate_answer(query, context)
        print(f"\nBot: {answer}\n")

if __name__ == "__main__":
    chat()

Run it with python chatbot.py. The low temperature of 0.1 keeps answers factual, and the strict system prompt forces the model to stay grounded in retrieved context.

Step 4: Adding a Web API

For production use, expose the chatbot as an HTTP endpoint. Create api.py:

from fastapi import FastAPI
from pydantic import BaseModel
from chatbot import load_vectorstore, retrieve_context, generate_answer

app = FastAPI(title="KB Chatbot API")
vectorstore = load_vectorstore()

class QueryRequest(BaseModel):
    question: str
    top_k: int = 4

class QueryResponse(BaseModel):
    answer: str
    context: str

@app.post("/ask", response_model=QueryResponse)
def ask(req: QueryRequest):
    context = retrieve_context(vectorstore, req.question, k=req.top_k)
    answer = generate_answer(req.question, context)
    return QueryResponse(answer=answer, context=context)

@app.get("/health")
def health():
    return {"status": "ok"}

Launch with uvicorn api:app --host 0.0.0.0 --port 9000. You now have a REST endpoint that any frontend can call.

Best Practices

Chunking Strategy

Chunk size is the single most impactful retrieval parameter. For technical documentation, 400-600 tokens with 10-15% overlap works well. For narrative documents, larger chunks of 800-1000 tokens preserve coherence. Experiment with your specific corpus and measure retrieval recall on a small evaluation set.

Hybrid Retrieval

Pure semantic search misses exact keyword matches. Combine FAISS similarity search with BM25 keyword search using a reciprocal rank fusion algorithm. This catches cases where the user uses exact product names, error codes, or identifiers that embeddings may not capture precisely.

Reranking

Retrieve a larger candidate set (k=20) and rerank with a cross-encoder model like cross-encoder/ms-marco-MiniLM-L-6-v2. Cross-encoders score query-document pairs jointly and are far more accurate than bi-encoder embeddings, at the cost of speed. Reranking the top 20 down to 4 typically improves answer quality substantially.

Context Window Management

Monitor total prompt length. With Mistral-7B's 8192 token context window and a 512-token system prompt, you have roughly 7000 tokens for context plus the user question. If your chunks are 512 tokens each and you retrieve 4, that is 2048 tokens of context, leaving comfortable headroom. Use vLLM's --max-model-len to match your actual needs and save memory.

Quantization for Deployment

If GPU memory is tight, serve a quantized model. vLLM supports AWQ and GPTQ checkpoints out of the box:

python -m vllm.entrypoints.openai.api_server \
  --model TheBloke/Mistral-7B-Instruct-v0.2-AWQ \
  --quantization awq \
  --port 8000 \
  --max-model-len 8192

An AWQ-quantized 7B model uses roughly 5GB of VRAM, making it feasible to run on consumer GPUs like the RTX 4060.

Streaming Responses

For better UX, enable streaming so tokens appear as they generate. Set "stream": true in the vLLM request payload and use Server-Sent Events on the FastAPI side. vLLM supports streaming natively through its OpenAI-compatible endpoint.

Evaluation

Build a small golden dataset of 30-50 question-answer pairs from your knowledge base. Measure retrieval recall (did the correct chunk appear in top-k?) and generation faithfulness (does the answer match the ground truth?). Track these metrics as you tune chunk size, top-k, temperature, and reranking thresholds.

Security Considerations

Sanitize retrieved content before injecting it into prompts to prevent prompt injection attacks. If your knowledge base contains user-generated content, a malicious document could instruct the model to ignore its system prompt. Consider wrapping retrieved text in delimiters and adding explicit instructions to treat context as data, not instructions.

Conclusion

Building a knowledge base chatbot with vLLM gives you a fast, grounded, and cost-effective AI assistant that answers from your own documents. The combination of PagedAttention-powered inference, semantic retrieval, and a strict grounding prompt produces responses that are both accurate and quick. Start with the basic pipeline above, then layer in reranking, hybrid search, streaming, and evaluation as your requirements grow. With vLLM handling the heavy lifting of inference, you can focus on what matters most: the quality of your retrieval pipeline and the clarity of your prompts.

— Ad —

Google AdSense will appear here after approval

← Back to all articles