Building a RAG Chatbot — Create Your Own Document QA Bot with LangChain + ChromaDB + OpenAI

Documents → Chunking → Embedding → Store in Vector DB

Question → Embedding → Similar Document Search → Prompt Construction → LLM Answer

pip install langchain langchain-openai langchain-community \
            chromadb pypdf tiktoken streamlit python-dotenv

# .env
OPENAI_API_KEY=sk-proj-your-api-key-here

# config.py
import os
from dotenv import load_dotenv

load_dotenv()

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
CHROMA_PERSIST_DIR = "./chroma_db"
COLLECTION_NAME = "my_documents"
CHUNK_SIZE = 1000
CHUNK_OVERLAP = 200
EMBEDDING_MODEL = "text-embedding-3-small"
LLM_MODEL = "gpt-4o-mini"

# document_loader.py
from langchain_community.document_loaders import (
    PyPDFLoader,
    DirectoryLoader,
    TextLoader,
)

def load_pdf(file_path: str):
    """Load a single PDF file"""
    loader = PyPDFLoader(file_path)
    documents = loader.load()
    print(f"Loaded {len(documents)} pages from {file_path}")
    return documents

def load_directory(dir_path: str, glob: str = "**/*.pdf"):
    """Load all PDFs in a directory"""
    loader = DirectoryLoader(
        dir_path,
        glob=glob,
        loader_cls=PyPDFLoader,
        show_progress=True,
    )
    documents = loader.load()
    print(f"Loaded {len(documents)} pages from {dir_path}")
    return documents

# Usage example
documents = load_directory("./docs")

# chunker.py
from langchain.text_splitter import RecursiveCharacterTextSplitter

def chunk_documents(documents, chunk_size=1000, chunk_overlap=200):
    """Split documents into chunks"""
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        length_function=len,
        separators=["\n\n", "\n", ".", " ", ""],
    )

    chunks = text_splitter.split_documents(documents)
    print(f"Split {len(documents)} documents into {len(chunks)} chunks")

    # Add chunk index to metadata
    for i, chunk in enumerate(chunks):
        chunk.metadata["chunk_index"] = i
        chunk.metadata["chunk_size"] = len(chunk.page_content)

    return chunks

chunks = chunk_documents(documents)

# vectorstore.py
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
from config import CHROMA_PERSIST_DIR, COLLECTION_NAME, EMBEDDING_MODEL

def create_vectorstore(chunks):
    """Embed chunks and store in ChromaDB"""
    embeddings = OpenAIEmbeddings(model=EMBEDDING_MODEL)

    vectorstore = Chroma.from_documents(
        documents=chunks,
        embedding=embeddings,
        persist_directory=CHROMA_PERSIST_DIR,
        collection_name=COLLECTION_NAME,
    )

    print(f"Stored {len(chunks)} chunks in ChromaDB")
    return vectorstore

def load_vectorstore():
    """Load existing ChromaDB"""
    embeddings = OpenAIEmbeddings(model=EMBEDDING_MODEL)

    vectorstore = Chroma(
        persist_directory=CHROMA_PERSIST_DIR,
        collection_name=COLLECTION_NAME,
        embedding_function=embeddings,
    )

    count = vectorstore._collection.count()
    print(f"Loaded ChromaDB with {count} documents")
    return vectorstore

# retriever.py

def get_retriever(vectorstore, search_type="mmr", k=4):
    """Create a retriever from the vector store"""

    if search_type == "mmr":
        # MMR: Balance between relevance and diversity
        retriever = vectorstore.as_retriever(
            search_type="mmr",
            search_kwargs={
                "k": k,
                "fetch_k": 20,      # Number of candidate documents
                "lambda_mult": 0.7,  # 1.0=relevance, 0.0=diversity
            },
        )
    elif search_type == "similarity_score":
        # Similarity threshold-based
        retriever = vectorstore.as_retriever(
            search_type="similarity_score_threshold",
            search_kwargs={
                "score_threshold": 0.7,
                "k": k,
            },
        )
    else:
        # Default similarity search
        retriever = vectorstore.as_retriever(
            search_kwargs={"k": k},
        )

    return retriever

# rag_chain.py
from langchain_openai import ChatOpenAI
from langchain.prompts import ChatPromptTemplate
from langchain.schema.runnable import RunnablePassthrough
from langchain.schema.output_parser import StrOutputParser
from config import LLM_MODEL

SYSTEM_TEMPLATE = """You are a document-based QA assistant.
Answer the question using only the provided context.

Rules:
1. If the information is not in the context, reply "I could not find that information in the documents."
2. Write your answer in Korean.
3. Include specific numbers or citations when possible.
4. Mention the source documents that support your answer.

Context:
{context}
"""

def format_docs(docs):
    """Format retrieved documents"""
    formatted = []
    for i, doc in enumerate(docs, 1):
        source = doc.metadata.get("source", "unknown")
        page = doc.metadata.get("page", "?")
        formatted.append(
            f"[Document {i}] (Source: {source}, Page: {page})\n{doc.page_content}"
        )
    return "\n\n---\n\n".join(formatted)

def create_rag_chain(retriever):
    """Create a RAG chain"""
    llm = ChatOpenAI(
        model=LLM_MODEL,
        temperature=0,
        max_tokens=2000,
    )

    prompt = ChatPromptTemplate.from_messages([
        ("system", SYSTEM_TEMPLATE),
        ("human", "{question}"),
    ])

    chain = (
        {
            "context": retriever | format_docs,
            "question": RunnablePassthrough(),
        }
        | prompt
        | llm
        | StrOutputParser()
    )

    return chain

# conversation.py
from langchain.memory import ConversationBufferWindowMemory
from langchain.chains import ConversationalRetrievalChain
from langchain_openai import ChatOpenAI
from config import LLM_MODEL

def create_conversational_chain(retriever):
    """Create a RAG chain with conversation history support"""
    llm = ChatOpenAI(model=LLM_MODEL, temperature=0)

    memory = ConversationBufferWindowMemory(
        memory_key="chat_history",
        return_messages=True,
        output_key="answer",
        k=5,  # Keep only the last 5 turns
    )

    chain = ConversationalRetrievalChain.from_llm(
        llm=llm,
        retriever=retriever,
        memory=memory,
        return_source_documents=True,
        verbose=False,
    )

    return chain

# Usage example
chain = create_conversational_chain(retriever)

result = chain.invoke({"question": "What is the key content of this document?"})
print(result["answer"])
print(f"\nReferenced {len(result['source_documents'])} documents")

# app.py
import streamlit as st
from document_loader import load_pdf
from chunker import chunk_documents
from vectorstore import create_vectorstore, load_vectorstore
from retriever import get_retriever
from rag_chain import create_rag_chain

st.set_page_config(page_title="📚 Document QA Chatbot", layout="wide")
st.title("📚 RAG Document QA Chatbot")

# Sidebar: Document Upload
with st.sidebar:
    st.header("📁 Upload Documents")
    uploaded_files = st.file_uploader(
        "Upload PDF files",
        type=["pdf"],
        accept_multiple_files=True,
    )

    if uploaded_files and st.button("🔄 Process Documents"):
        with st.spinner("Processing documents..."):
            all_chunks = []
            for file in uploaded_files:
                # Save temporary file
                temp_path = f"/tmp/{file.name}"
                with open(temp_path, "wb") as f:
                    f.write(file.getbuffer())

                docs = load_pdf(temp_path)
                chunks = chunk_documents(docs)
                all_chunks.extend(chunks)

            vectorstore = create_vectorstore(all_chunks)
            st.session_state["vectorstore"] = vectorstore
            st.success(f"✅ {len(all_chunks)} chunks processed!")

# Main: Chat Interface
if "messages" not in st.session_state:
    st.session_state.messages = []

# Display previous messages
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# User input
if prompt := st.chat_input("Ask a question about your documents..."):
    st.session_state.messages.append({"role": "user", "content": prompt})

    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        if "vectorstore" not in st.session_state:
            try:
                st.session_state["vectorstore"] = load_vectorstore()
            except Exception:
                st.error("Please upload documents first.")
                st.stop()

        vectorstore = st.session_state["vectorstore"]
        retriever = get_retriever(vectorstore)
        chain = create_rag_chain(retriever)

        with st.spinner("Generating answer..."):
            response = chain.invoke(prompt)

        st.markdown(response)
        st.session_state.messages.append(
            {"role": "assistant", "content": response}
        )

# Run
streamlit run app.py --server.port 8501

from langchain.retrievers import EnsembleRetriever
from langchain_community.retrievers import BM25Retriever

def create_hybrid_retriever(chunks, vectorstore, k=4):
    """BM25 + Vector search ensemble"""
    bm25_retriever = BM25Retriever.from_documents(chunks)
    bm25_retriever.k = k

    vector_retriever = vectorstore.as_retriever(search_kwargs={"k": k})

    ensemble = EnsembleRetriever(
        retrievers=[bm25_retriever, vector_retriever],
        weights=[0.4, 0.6],  # Weight toward vector search
    )

    return ensemble

from langchain.retrievers import ContextualCompressionRetriever
from langchain_community.document_compressors import CohereRerank

def create_reranked_retriever(vectorstore, k=4, top_n=3):
    """Rerank search results using Cohere Reranker"""
    base_retriever = vectorstore.as_retriever(search_kwargs={"k": k * 3})

    compressor = CohereRerank(
        model="rerank-v3.5",
        top_n=top_n,
    )

    return ContextualCompressionRetriever(
        base_compressor=compressor,
        base_retriever=base_retriever,
    )

# Add summary metadata to chunks
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

for chunk in chunks:
    summary = llm.invoke(
        f"Summarize the following text in one sentence:\n{chunk.page_content}"
    ).content
    chunk.metadata["summary"] = summary