Deploying a Medical Chatbot with LlamaIndex and RAG

Introduction

Large Language Models (LLMs) like GPT-4 and Llama-3 are incredibly smart, but they are prone to hallucinations and do not have access to private hospital data or the very latest clinical guidelines. Retrieval-Augmented Generation (RAG) solves this by retrieving relevant documents from a private database and injecting them into the LLM's prompt. In this tutorial, we will build a medical chatbot using LlamaIndex to securely query local clinical PDF documents.

Architecture Overview

sequenceDiagram
    participant User as Doctor (User)
    participant Chat as LlamaIndex Chat Engine
    participant DB as Vector Database
    participant LLM as Language Model
    
    Note over Chat, DB: Indexing Phase (One-Time)
    Chat->>DB: Read & Chunk Clinical PDFs
    Chat->>LLM: Generate Vector Embeddings for Chunks
    LLM-->>DB: Store Embeddings
    
    Note over User, LLM: Query Phase
    User->>Chat: "What is the recommended dosage for drug X?"
    Chat->>LLM: Embed User Question
    LLM-->>Chat: Question Embedding
    Chat->>DB: Semantic Search (Cosine Similarity)
    DB-->>Chat: Return Top-K Relevant Document Chunks
    Chat->>LLM: Prompt: Answer question using ONLY these chunks
    LLM-->>User: "Based on the internal SOPs, the dosage is..."
        

Prerequisites

• Python 3.9+
• An OpenAI API Key (or a local LLM via Ollama)
• Basic understanding of vector embeddings

Step 1: Install Dependencies

pip install llama-index pypdf python-dotenv

Step 2: Preparing the Clinical Documents

Create a directory called clinical_docs/ and place some PDF files in it. For example, you might download the latest AHA/ACC cardiovascular guidelines or a hospital's internal Standard Operating Procedures (SOPs).

mkdir clinical_docs
# Place your PDFs inside clinical_docs/

Step 3: Loading and Indexing the Documents

LlamaIndex makes it incredibly easy to load documents, chunk them into smaller pieces, convert those chunks into vector embeddings (using OpenAI's embedding model), and store them in an in-memory vector database.

import os
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
from dotenv import load_dotenv

# Load your OpenAI API key from a .env file
load_dotenv()
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")

# Optional: Configure chunk size (default is 1024 tokens)
# Smaller chunks often yield more precise retrieval
Settings.chunk_size = 512
Settings.chunk_overlap = 50

print("Loading clinical documents...")
documents = SimpleDirectoryReader('clinical_docs').load_data()

print("Creating vector embeddings and indexing...")
index = VectorStoreIndex.from_documents(documents)

# Save the index to disk so we don't have to re-compute embeddings every time
index.storage_context.persist(persist_dir="./storage")
print("Index saved to ./storage")

Step 4: Querying the Chatbot

Now that the documents are indexed, we can create a chat engine. When the user asks a question, LlamaIndex will automatically search the vector database for the most relevant chunks of text, and provide those chunks to the LLM to formulate an answer.

from llama_index.core import StorageContext, load_index_from_storage

# Re-load the index from disk
storage_context = StorageContext.from_defaults(persist_dir="./storage")
index = load_index_from_storage(storage_context)

# Initialize the chat engine
# We use a custom system prompt to strictly enforce clinical safety
custom_prompt = (
    "You are a clinical AI assistant. You must answer questions based strictly "
    "on the provided context documents. If the context does not contain the answer, "
    "you must clearly state 'I do not know'. Do not hallucinate or use outside knowledge."
)

chat_engine = index.as_chat_engine(
    chat_mode="condense_question", 
    verbose=True,
    system_prompt=custom_prompt
)

print("Medical Chatbot Initialized. Type 'exit' to quit.")
while True:
    user_input = input("\nDoctor: ")
    if user_input.lower() == 'exit':
        break
    
    response = chat_engine.chat(user_input)
    print(f"AI Assistant: {response.response}")
    
    # Optional: Print the source documents that were used to generate the answer
    # for source in response.source_nodes:
    #     print(f"Source Node: {source.node.get_text()[:100]}...")

Troubleshooting: Common RAG Pitfalls

• Hallucinations Despite RAG: If the model is still making things up, ensure your system_prompt is sufficiently strict. You can also explicitly instruct it to cite its sources.
• "Lost in the Middle": If you retrieve too many chunks (e.g., top-k = 20), LLMs often ignore information in the middle of the prompt. Try reducing similarity_top_k or enabling a Re-ranker (like Cohere ReRank).
• Poor Retrieval Quality: Semantic search struggles with highly specific medical acronyms. Consider implementing a Hybrid Search (Semantic + BM25 Keyword Search) to catch exact matches for drug names or obscure genes.

Conclusion

With just a few lines of Python using LlamaIndex, we've built a powerful, context-aware clinical assistant. Unlike a raw LLM, this RAG architecture ensures the chatbot bases its answers strictly on the provided documents, dramatically reducing hallucinations—a critical requirement for any AI deployed in a healthcare setting.