RAG from Scratch: Building a Retrieval-Augmented Chatbot Over Your Own Documents

1
2
3
4
5
6

# Only search within documents tagged as "architecture"
results = vectorstore.similarity_search(
    query,
    k=5,
    filter={"category": "architecture"}
)

1
2
3
4
5
6

# Install Python packages
pip install langchain langchain-community langchain-ollama chromadb pypdf python-dotenv

# Pull the models you need via Ollama
ollama pull llama3.2
ollama pull nomic-embed-text

1
2

ollama list
# Should show llama3.2 and nomic-embed-text

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68

# loaders.py
from pathlib import Path
from langchain_community.document_loaders import (
    PyPDFLoader,
    DirectoryLoader,
    TextLoader,
    WebBaseLoader,
    UnstructuredMarkdownLoader,
)
from langchain.schema import Document
from typing import List


def load_pdfs(path: str) -> List[Document]:
    """Load all PDFs from a directory, one Document per page."""
    loader = DirectoryLoader(
        path,
        glob="**/*.pdf",
        loader_cls=PyPDFLoader,
        show_progress=True,
    )
    docs = loader.load()
    print(f"Loaded {len(docs)} pages from PDFs in {path}")
    return docs


def load_markdown(path: str) -> List[Document]:
    """Load all markdown files from a directory."""
    loader = DirectoryLoader(
        path,
        glob="**/*.md",
        loader_cls=UnstructuredMarkdownLoader,
        show_progress=True,
    )
    docs = loader.load()
    print(f"Loaded {len(docs)} markdown documents from {path}")
    return docs


def load_text_files(path: str) -> List[Document]:
    """Load all plain text files from a directory."""
    loader = DirectoryLoader(
        path,
        glob="**/*.txt",
        loader_cls=TextLoader,
        loader_kwargs={"encoding": "utf-8"},
        show_progress=True,
    )
    docs = loader.load()
    print(f"Loaded {len(docs)} text files from {path}")
    return docs


def load_web_page(url: str) -> List[Document]:
    """Load a web page and return its text content."""
    loader = WebBaseLoader(url)
    docs = loader.load()
    print(f"Loaded {len(docs)} documents from {url}")
    return docs


def load_all_documents(directory: str) -> List[Document]:
    """Load PDFs, markdown, and text files from a directory."""
    docs = []
    docs.extend(load_pdfs(directory))
    docs.extend(load_markdown(directory))
    docs.extend(load_text_files(directory))
    return docs

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84

# splitter.py
from langchain.text_splitter import (
    RecursiveCharacterTextSplitter,
    MarkdownHeaderTextSplitter,
    Language,
)
from langchain.schema import Document
from typing import List


def split_documents(
    docs: List[Document],
    chunk_size: int = 1000,
    chunk_overlap: int = 200,
) -> List[Document]:
    """
    Split documents into chunks using RecursiveCharacterTextSplitter.

    chunk_size: maximum number of characters per chunk
    chunk_overlap: number of characters shared between adjacent chunks

    The splitter tries to break on natural boundaries in this order:
    paragraph breaks (\n\n), line breaks (\n), sentences (. ), words ( ),
    and finally individual characters. It falls back to the next separator
    only when the current one would produce a chunk that exceeds chunk_size.
    """
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        length_function=len,
        separators=["\n\n", "\n", ". ", " ", ""],
    )

    chunks = splitter.split_documents(docs)
    print(f"Split {len(docs)} documents into {len(chunks)} chunks")
    print(f"Average chunk size: {sum(len(c.page_content) for c in chunks) // len(chunks)} chars")
    return chunks


def split_markdown_by_headers(markdown_text: str) -> List[Document]:
    """
    Split a markdown document by its header hierarchy.
    Produces chunks that align with the document's logical structure.
    """
    headers_to_split_on = [
        ("#", "h1"),
        ("##", "h2"),
        ("###", "h3"),
    ]
    header_splitter = MarkdownHeaderTextSplitter(
        headers_to_split_on=headers_to_split_on,
        strip_headers=False,
    )
    sections = header_splitter.split_text(markdown_text)

    # Further split large sections so no chunk exceeds our size limit
    char_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=100,
    )
    return char_splitter.split_documents(sections)


def split_code(code_text: str, language: str = "python") -> List[Document]:
    """
    Split source code using LangChain's language-aware splitter.
    Breaks on class and function boundaries rather than arbitrary character counts.
    """
    lang_map = {
        "python": Language.PYTHON,
        "javascript": Language.JS,
        "typescript": Language.TS,
        "go": Language.GO,
        "rust": Language.RUST,
        "java": Language.JAVA,
    }
    lang = lang_map.get(language.lower(), Language.PYTHON)

    splitter = RecursiveCharacterTextSplitter.from_language(
        language=lang,
        chunk_size=1500,
        chunk_overlap=200,
    )
    return splitter.create_documents([code_text])

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61

# indexer.py
from langchain_ollama import OllamaEmbeddings
from langchain_community.vectorstores import Chroma
from langchain.schema import Document
from typing import List
import os


CHROMA_PERSIST_DIR = "./chroma_db"
EMBEDDING_MODEL = "nomic-embed-text"
COLLECTION_NAME = "my_documents"


def get_embeddings() -> OllamaEmbeddings:
    """Return the embedding model. Must be the same at index and query time."""
    return OllamaEmbeddings(model=EMBEDDING_MODEL)


def build_index(chunks: List[Document]) -> Chroma:
    """
    Embed all chunks and store them in ChromaDB.
    ChromaDB persists to disk automatically at CHROMA_PERSIST_DIR.
    """
    embeddings = get_embeddings()

    # If the collection already exists, delete and recreate it
    # (for a full re-index — see incremental updates section for smarter logic)
    vectorstore = Chroma.from_documents(
        documents=chunks,
        embedding=embeddings,
        collection_name=COLLECTION_NAME,
        persist_directory=CHROMA_PERSIST_DIR,
    )

    print(f"Indexed {len(chunks)} chunks into ChromaDB at {CHROMA_PERSIST_DIR}")
    return vectorstore


def load_index() -> Chroma:
    """Load an existing ChromaDB index from disk."""
    if not os.path.exists(CHROMA_PERSIST_DIR):
        raise FileNotFoundError(
            f"No ChromaDB index found at {CHROMA_PERSIST_DIR}. "
            "Run the indexing pipeline first."
        )
    embeddings = get_embeddings()
    vectorstore = Chroma(
        collection_name=COLLECTION_NAME,
        embedding_function=embeddings,
        persist_directory=CHROMA_PERSIST_DIR,
    )
    count = vectorstore._collection.count()
    print(f"Loaded ChromaDB index with {count} chunks")
    return vectorstore


def add_documents_to_index(new_chunks: List[Document]) -> None:
    """Add new chunks to an existing index without rebuilding from scratch."""
    vectorstore = load_index()
    vectorstore.add_documents(new_chunks)
    print(f"Added {len(new_chunks)} new chunks to the index")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74

# retriever.py
from langchain_community.vectorstores import Chroma
from langchain.schema import Document
from typing import List, Tuple
from indexer import load_index


def similarity_search(
    query: str,
    k: int = 5,
    score_threshold: float = 0.3,
) -> List[Tuple[Document, float]]:
    """
    Return the top-k most similar chunks to the query, with their scores.
    score_threshold filters out chunks below a minimum similarity.
    """
    vectorstore = load_index()
    results = vectorstore.similarity_search_with_relevance_scores(
        query,
        k=k,
    )

    # Filter by score threshold
    filtered = [(doc, score) for doc, score in results if score >= score_threshold]

    if not filtered:
        print(f"Warning: no chunks above score threshold {score_threshold}")
        return results[:k]  # Return top-k anyway

    return filtered


def mmr_search(
    query: str,
    k: int = 5,
    fetch_k: int = 20,
    lambda_mult: float = 0.5,
) -> List[Document]:
    """
    Maximal Marginal Relevance search.

    Retrieves fetch_k candidates by similarity, then selects k of them
    to maximize both relevance AND diversity (penalizes redundant chunks).

    lambda_mult: 0.0 = pure diversity, 1.0 = pure similarity.
    0.5 is a balanced default.
    """
    vectorstore = load_index()
    return vectorstore.max_marginal_relevance_search(
        query,
        k=k,
        fetch_k=fetch_k,
        lambda_mult=lambda_mult,
    )


def filtered_search(
    query: str,
    metadata_filter: dict,
    k: int = 5,
) -> List[Document]:
    """
    Restrict retrieval to chunks matching a metadata filter.

    Example:
        filtered_search("nginx config", {"source": "runbooks/nginx.md"})
        filtered_search("backup policy", {"doc_type": "policy"})
    """
    vectorstore = load_index()
    return vectorstore.similarity_search(
        query,
        k=k,
        filter=metadata_filter,
    )

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103

# rag_chain.py
from langchain_ollama import OllamaLLM
from langchain.prompts import ChatPromptTemplate
from langchain.schema import Document
from retriever import similarity_search, mmr_search
from typing import List, Dict, Any


LLM_MODEL = "llama3.2"

SYSTEM_PROMPT = """You are a helpful assistant that answers questions based strictly on the provided context documents.

Rules:
- Answer only using information present in the context below.
- If the context does not contain enough information to answer the question, say so explicitly — do not guess or hallucinate.
- When referencing specific information, mention which document it came from.
- Be concise and precise. Do not add information that is not in the context.
- If the question is ambiguous, answer for the most likely interpretation and note the ambiguity.

Context:
{context}
"""

HUMAN_PROMPT = "Question: {question}"


def format_docs(docs_with_scores: List[tuple]) -> str:
    """
    Format retrieved documents into a string for injection into the prompt.
    Includes source metadata so the model can reference it.
    """
    formatted = []
    for i, (doc, score) in enumerate(docs_with_scores, 1):
        source = doc.metadata.get("source", "unknown")
        page = doc.metadata.get("page", "")
        page_str = f", page {page}" if page else ""
        formatted.append(
            f"[Document {i} | Source: {source}{page_str} | Relevance: {score:.2f}]\n"
            f"{doc.page_content}"
        )
    return "\n\n---\n\n".join(formatted)


def get_llm() -> OllamaLLM:
    return OllamaLLM(model=LLM_MODEL, temperature=0.1)


def ask(
    question: str,
    k: int = 5,
    use_mmr: bool = False,
    metadata_filter: dict = None,
) -> Dict[str, Any]:
    """
    Full RAG pipeline: retrieve relevant chunks, build prompt, generate answer.

    Returns a dict with:
        - answer: the LLM's response
        - sources: list of (document, score) tuples
        - context: the formatted context string sent to the LLM
    """
    # Retrieve
    if use_mmr:
        docs = mmr_search(question, k=k)
        docs_with_scores = [(doc, 0.0) for doc in docs]  # MMR doesn't return scores
    else:
        docs_with_scores = similarity_search(question, k=k)
        if metadata_filter:
            from retriever import filtered_search
            filtered_docs = filtered_search(question, metadata_filter, k=k)
            docs_with_scores = [(doc, 0.0) for doc in filtered_docs]

    if not docs_with_scores:
        return {
            "answer": "I could not find any relevant documents to answer your question.",
            "sources": [],
            "context": "",
        }

    # Format context
    context = format_docs(docs_with_scores)

    # Build prompt
    prompt = ChatPromptTemplate.from_messages([
        ("system", SYSTEM_PROMPT),
        ("human", HUMAN_PROMPT),
    ])

    # Build chain
    llm = get_llm()
    chain = prompt | llm

    # Generate
    answer = chain.invoke({
        "context": context,
        "question": question,
    })

    return {
        "answer": answer,
        "sources": docs_with_scores,
        "context": context,
    }

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251

#!/usr/bin/env python3
"""
rag.py — Complete RAG pipeline in a single script.

Usage:
    # Index a directory of documents
    python rag.py index ./my_docs

    # Ask a question
    python rag.py ask "How do I configure nginx for SSL termination?"

    # Ask with MMR retrieval (better diversity)
    python rag.py ask "What are the backup procedures?" --mmr

    # Ask with metadata filter
    python rag.py ask "database schema" --filter '{"doc_type": "technical"}'
"""

import sys
import json
import argparse
from pathlib import Path
from typing import List

from langchain_community.document_loaders import (
    PyPDFLoader,
    DirectoryLoader,
    TextLoader,
    UnstructuredMarkdownLoader,
)
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_ollama import OllamaEmbeddings, OllamaLLM
from langchain_community.vectorstores import Chroma
from langchain.prompts import ChatPromptTemplate
from langchain.schema import Document


# --- Configuration ---
CHROMA_PERSIST_DIR = "./chroma_db"
COLLECTION_NAME = "my_documents"
EMBEDDING_MODEL = "nomic-embed-text"
LLM_MODEL = "llama3.2"
CHUNK_SIZE = 1000
CHUNK_OVERLAP = 200

SYSTEM_PROMPT = """You are a helpful assistant that answers questions based strictly on the provided context documents.

Rules:
- Answer only using information present in the context below.
- If the context does not contain enough information to answer the question, say so explicitly.
- When referencing specific information, mention which document it came from.
- Be concise and precise. Do not add information that is not in the context.

Context:
{context}
"""


# --- Indexing ---

def load_documents(directory: str) -> List[Document]:
    docs = []
    path = Path(directory)

    if not path.exists():
        raise FileNotFoundError(f"Directory not found: {directory}")

    loaders = [
        ("**/*.pdf", PyPDFLoader),
        ("**/*.md", UnstructuredMarkdownLoader),
        ("**/*.txt", TextLoader),
    ]

    for glob_pattern, loader_cls in loaders:
        if list(path.glob(glob_pattern)):
            kwargs = {}
            if loader_cls == TextLoader:
                kwargs["loader_kwargs"] = {"encoding": "utf-8"}
            loader = DirectoryLoader(
                directory,
                glob=glob_pattern,
                loader_cls=loader_cls,
                show_progress=True,
                **kwargs,
            )
            loaded = loader.load()
            docs.extend(loaded)
            print(f"  Loaded {len(loaded)} documents matching {glob_pattern}")

    return docs


def split_documents(docs: List[Document]) -> List[Document]:
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=CHUNK_SIZE,
        chunk_overlap=CHUNK_OVERLAP,
        separators=["\n\n", "\n", ". ", " ", ""],
    )
    chunks = splitter.split_documents(docs)
    return chunks


def index_documents(directory: str) -> None:
    print(f"\n=== Indexing documents from: {directory} ===\n")

    print("Loading documents...")
    docs = load_documents(directory)
    if not docs:
        print("No documents found. Check the directory and supported formats (.pdf, .md, .txt).")
        sys.exit(1)
    print(f"Total documents loaded: {len(docs)}")

    print("\nSplitting into chunks...")
    chunks = split_documents(docs)
    avg_size = sum(len(c.page_content) for c in chunks) // len(chunks)
    print(f"Total chunks: {len(chunks)} | Average chunk size: {avg_size} chars")

    print(f"\nEmbedding and storing in ChromaDB (model: {EMBEDDING_MODEL})...")
    print("(This may take a few minutes for large document sets.)\n")

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

    count = vectorstore._collection.count()
    print(f"\nIndexing complete. {count} chunks stored at {CHROMA_PERSIST_DIR}")


# --- Querying ---

def load_vectorstore() -> Chroma:
    embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
    return Chroma(
        collection_name=COLLECTION_NAME,
        embedding_function=embeddings,
        persist_directory=CHROMA_PERSIST_DIR,
    )


def format_docs_for_prompt(docs_with_scores) -> str:
    parts = []
    for i, (doc, score) in enumerate(docs_with_scores, 1):
        source = doc.metadata.get("source", "unknown")
        page = doc.metadata.get("page", "")
        loc = f"{source}" + (f" (page {page})" if page else "")
        parts.append(
            f"[Chunk {i} | {loc} | score={score:.3f}]\n{doc.page_content}"
        )
    return "\n\n---\n\n".join(parts)


def answer_question(
    question: str,
    k: int = 5,
    use_mmr: bool = False,
    metadata_filter: dict = None,
) -> None:
    print(f"\n=== Question: {question} ===\n")

    vectorstore = load_vectorstore()

    # Retrieve
    if use_mmr:
        raw_docs = vectorstore.max_marginal_relevance_search(
            question, k=k, fetch_k=k * 4, lambda_mult=0.5
        )
        docs_with_scores = [(doc, 0.0) for doc in raw_docs]
    elif metadata_filter:
        raw_docs = vectorstore.similarity_search(
            question, k=k, filter=metadata_filter
        )
        docs_with_scores = [(doc, 0.0) for doc in raw_docs]
    else:
        docs_with_scores = vectorstore.similarity_search_with_relevance_scores(
            question, k=k
        )

    if not docs_with_scores:
        print("No relevant documents found.")
        return

    # Show retrieved chunks
    print(f"Retrieved {len(docs_with_scores)} chunks:\n")
    for i, (doc, score) in enumerate(docs_with_scores, 1):
        source = doc.metadata.get("source", "unknown")
        preview = doc.page_content[:120].replace("\n", " ")
        print(f"  [{i}] score={score:.3f} | {source}")
        print(f"      \"{preview}...\"")
    print()

    # Build prompt and generate
    context = format_docs_for_prompt(docs_with_scores)
    prompt = ChatPromptTemplate.from_messages([
        ("system", SYSTEM_PROMPT),
        ("human", "Question: {question}"),
    ])
    llm = OllamaLLM(model=LLM_MODEL, temperature=0.1)
    chain = prompt | llm

    print("=== Answer ===\n")
    answer = chain.invoke({"context": context, "question": question})
    print(answer)

    print("\n=== Sources ===")
    seen = set()
    for doc, score in docs_with_scores:
        source = doc.metadata.get("source", "unknown")
        if source not in seen:
            seen.add(source)
            print(f"  - {source}")


# --- CLI ---

def main():
    parser = argparse.ArgumentParser(description="RAG from Scratch")
    subparsers = parser.add_subparsers(dest="command")

    # Index command
    idx = subparsers.add_parser("index", help="Index a directory of documents")
    idx.add_argument("directory", help="Path to the document directory")

    # Ask command
    ask_p = subparsers.add_parser("ask", help="Ask a question")
    ask_p.add_argument("question", help="The question to answer")
    ask_p.add_argument("--k", type=int, default=5, help="Number of chunks to retrieve")
    ask_p.add_argument("--mmr", action="store_true", help="Use MMR for diverse retrieval")
    ask_p.add_argument("--filter", type=str, default=None, help="JSON metadata filter")

    args = parser.parse_args()

    if args.command == "index":
        index_documents(args.directory)
    elif args.command == "ask":
        metadata_filter = json.loads(args.filter) if args.filter else None
        answer_question(
            args.question,
            k=args.k,
            use_mmr=args.mmr,
            metadata_filter=metadata_filter,
        )
    else:
        parser.print_help()


if __name__ == "__main__":
    main()

1
2
3
4
5
6
7

# Index your documents
python rag.py index ./docs

# Ask questions
python rag.py ask "What is our incident response procedure?"
python rag.py ask "How do I deploy to production?" --mmr
python rag.py ask "nginx configuration" --filter '{"source": "runbooks/nginx.md"}'

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

from langchain_ollama import OllamaLLM
from langchain.prompts import PromptTemplate

def expand_query(question: str, n: int = 3) -> list[str]:
    """Generate n alternative phrasings of a question for improved recall."""
    llm = OllamaLLM(model="llama3.2", temperature=0.7)
    prompt = PromptTemplate.from_template(
        "Generate {n} different ways to phrase this question for a document search. "
        "Output only the alternative questions, one per line, with no numbering or explanation.\n\n"
        "Original question: {question}"
    )
    result = llm.invoke(prompt.format(n=n, question=question))
    alternatives = [q.strip() for q in result.strip().split("\n") if q.strip()]
    return [question] + alternatives[:n]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

def hyde_search(question: str, k: int = 5) -> list:
    """
    Hypothetical Document Embeddings search.
    Generate a hypothetical answer, embed it, search with that embedding.
    """
    llm = OllamaLLM(model="llama3.2", temperature=0.5)
    prompt = PromptTemplate.from_template(
        "Write a detailed technical paragraph that answers this question. "
        "Write as if you are a technical document, not as if you are directly answering. "
        "If you don't know the answer, make a plausible technical response anyway.\n\n"
        "Question: {question}\n\nParagraph:"
    )
    hypothetical_doc = llm.invoke(prompt.format(question=question))

    vectorstore = load_vectorstore()
    return vectorstore.similarity_search(hypothetical_doc, k=k)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21

from sentence_transformers import CrossEncoder

def rerank_results(query: str, docs: list, top_k: int = 5) -> list:
    """
    Re-rank retrieved documents using a cross-encoder.
    More accurate than embedding similarity alone.
    """
    model = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")

    # Score each (query, document) pair
    pairs = [[query, doc.page_content] for doc in docs]
    scores = model.predict(pairs)

    # Sort by score descending
    scored = sorted(zip(docs, scores), key=lambda x: x[1], reverse=True)
    return [doc for doc, score in scored[:top_k]]


# Usage: retrieve 20 candidates, re-rank to top 5
candidates = vectorstore.similarity_search(question, k=20)
reranked = rerank_results(question, candidates, top_k=5)

1

pip install sentence-transformers

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144

# app.py
import streamlit as st
from pathlib import Path
import tempfile
import shutil
import os

from langchain_community.document_loaders import PyPDFLoader, TextLoader, UnstructuredMarkdownLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_ollama import OllamaEmbeddings, OllamaLLM
from langchain_community.vectorstores import Chroma
from langchain.prompts import ChatPromptTemplate

# --- Config ---
EMBEDDING_MODEL = "nomic-embed-text"
LLM_MODEL = "llama3.2"
CHROMA_DIR = "./streamlit_chroma_db"
CHUNK_SIZE = 1000
CHUNK_OVERLAP = 200

SYSTEM_PROMPT = """You are a helpful assistant. Answer questions based only on the provided context.
If the context does not contain the answer, say so clearly. Do not fabricate information.

Context:
{context}"""


def index_uploaded_files(uploaded_files) -> Chroma:
    """Process uploaded files and build a vector index."""
    docs = []
    with tempfile.TemporaryDirectory() as tmpdir:
        for f in uploaded_files:
            file_path = os.path.join(tmpdir, f.name)
            with open(file_path, "wb") as out:
                out.write(f.getbuffer())

            if f.name.endswith(".pdf"):
                loader = PyPDFLoader(file_path)
            elif f.name.endswith(".md"):
                loader = UnstructuredMarkdownLoader(file_path)
            else:
                loader = TextLoader(file_path, encoding="utf-8")

            docs.extend(loader.load())

    splitter = RecursiveCharacterTextSplitter(
        chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP
    )
    chunks = splitter.split_documents(docs)

    embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
    # Clear and rebuild index
    if os.path.exists(CHROMA_DIR):
        shutil.rmtree(CHROMA_DIR)

    vectorstore = Chroma.from_documents(
        documents=chunks,
        embedding=embeddings,
        persist_directory=CHROMA_DIR,
    )
    return vectorstore


def get_answer(question: str, vectorstore: Chroma) -> dict:
    """Run the RAG pipeline for a single question."""
    docs_with_scores = vectorstore.similarity_search_with_relevance_scores(question, k=5)

    context_parts = []
    for i, (doc, score) in enumerate(docs_with_scores, 1):
        source = doc.metadata.get("source", "uploaded file")
        context_parts.append(f"[{i}] (source: {source})\n{doc.page_content}")
    context = "\n\n---\n\n".join(context_parts)

    prompt = ChatPromptTemplate.from_messages([
        ("system", SYSTEM_PROMPT),
        ("human", "{question}"),
    ])
    llm = OllamaLLM(model=LLM_MODEL, temperature=0.1)
    chain = prompt | llm

    answer = chain.invoke({"context": context, "question": question})
    sources = list({doc.metadata.get("source", "unknown") for doc, _ in docs_with_scores})
    return {"answer": answer, "sources": sources}


# --- Streamlit UI ---

st.set_page_config(page_title="RAG Chatbot", page_icon="", layout="wide")
st.title("Document Q&A — RAG Chatbot")

with st.sidebar:
    st.header("Upload Documents")
    uploaded_files = st.file_uploader(
        "Upload PDF, Markdown, or Text files",
        type=["pdf", "md", "txt"],
        accept_multiple_files=True,
    )

    if uploaded_files and st.button("Index Documents", type="primary"):
        with st.spinner(f"Indexing {len(uploaded_files)} file(s)..."):
            st.session_state.vectorstore = index_uploaded_files(uploaded_files)
            st.session_state.indexed_files = [f.name for f in uploaded_files]
        st.success(f"Indexed {len(uploaded_files)} file(s)")

    if "indexed_files" in st.session_state:
        st.subheader("Indexed files:")
        for fname in st.session_state.indexed_files:
            st.write(f"- {fname}")

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

for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        st.write(msg["content"])
        if msg.get("sources"):
            with st.expander("Sources"):
                for s in msg["sources"]:
                    st.write(f"- {s}")

if prompt := st.chat_input("Ask a question about your documents..."):
    if "vectorstore" not in st.session_state:
        st.error("Please upload and index documents first.")
    else:
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.write(prompt)

        with st.chat_message("assistant"):
            with st.spinner("Thinking..."):
                result = get_answer(prompt, st.session_state.vectorstore)

            st.write(result["answer"])
            if result["sources"]:
                with st.expander("Sources"):
                    for s in result["sources"]:
                        st.write(f"- {s}")

        st.session_state.messages.append({
            "role": "assistant",
            "content": result["answer"],
            "sources": result["sources"],
        })

1
2

pip install streamlit
streamlit run app.py

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

# docker-compose.yml
services:
  ollama:
    image: ollama/ollama:latest
    ports:
      - "11434:11434"
    volumes:
      - ollama_data:/root/.ollama
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]
    # Remove the deploy block if you don't have a GPU

  rag-app:
    build: .
    ports:
      - "8501:8501"
    environment:
      - OLLAMA_BASE_URL=http://ollama:11434
    depends_on:
      - ollama
    volumes:
      - ./chroma_db:/app/chroma_db
      - ./docs:/app/docs

volumes:
  ollama_data:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

# Dockerfile
FROM python:3.11-slim

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

COPY . .
EXPOSE 8501
CMD ["streamlit", "run", "app.py", "--server.address", "0.0.0.0"]

1

pip install ragas

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51

# evaluate.py
from ragas import evaluate
from ragas.metrics import (
    faithfulness,
    answer_relevancy,
    context_precision,
    context_recall,
)
from datasets import Dataset


def evaluate_rag_pipeline(qa_pairs: list[dict]) -> dict:
    """
    Evaluate a RAG pipeline on a set of question-answer pairs.

    qa_pairs: list of dicts with keys:
        - question: str
        - answer: str (the RAG system's answer)
        - contexts: list[str] (the retrieved chunks)
        - ground_truth: str (the correct answer, for recall calculation)
    """
    dataset = Dataset.from_list(qa_pairs)

    results = evaluate(
        dataset,
        metrics=[
            faithfulness,
            answer_relevancy,
            context_precision,
            context_recall,
        ],
    )
    return results


# Example usage
sample_qa = [
    {
        "question": "What is the maximum upload file size?",
        "answer": "The maximum upload file size is 50MB as configured in nginx.conf.",
        "contexts": [
            "The nginx configuration sets client_max_body_size to 50m, limiting uploads to 50MB.",
            "Files larger than 50MB will receive a 413 Request Entity Too Large error.",
        ],
        "ground_truth": "The maximum file upload size is 50MB.",
    },
]

scores = evaluate_rag_pipeline(sample_qa)
print(scores)
# Output: {'faithfulness': 0.97, 'answer_relevancy': 0.95, ...}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

def judge_answer(question: str, answer: str, context: str) -> dict:
    """Use an LLM to evaluate RAG output quality."""
    llm = OllamaLLM(model="llama3.2", temperature=0.0)
    prompt = f"""Evaluate this RAG system response on three dimensions.
Score each 1-5 (5 = best). Return JSON only.

Question: {question}

Retrieved Context:
{context}

System Answer: {answer}

Evaluate:
1. faithfulness (1-5): Is the answer fully supported by the context? Penalize any claims not in context.
2. relevance (1-5): Does the answer address the question?
3. completeness (1-5): Does the answer cover all relevant information from the context?

Return only valid JSON: {{"faithfulness": N, "relevance": N, "completeness": N, "reasoning": "brief note"}}"""

    result = llm.invoke(prompt)
    import json
    return json.loads(result.strip())

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

import hashlib
import json
from pathlib import Path


def get_file_hash(file_path: str) -> str:
    with open(file_path, "rb") as f:
        return hashlib.md5(f.read()).hexdigest()


def load_index_manifest(manifest_path: str = "./index_manifest.json") -> dict:
    if Path(manifest_path).exists():
        with open(manifest_path) as f:
            return json.load(f)
    return {}


def save_index_manifest(manifest: dict, manifest_path: str = "./index_manifest.json"):
    with open(manifest_path, "w") as f:
        json.dump(manifest, f, indent=2)


def get_changed_files(directory: str, manifest: dict) -> list[str]:
    """Return files that are new or changed since last index."""
    changed = []
    for path in Path(directory).rglob("*"):
        if path.is_file() and path.suffix in {".pdf", ".md", ".txt"}:
            file_hash = get_file_hash(str(path))
            if manifest.get(str(path)) != file_hash:
                changed.append(str(path))
    return changed

1
2
3
4
5
6
7
8
9

# When indexing, add user/group metadata
chunk.metadata["owner"] = "team-ops"
chunk.metadata["access_level"] = "internal"

# At query time, filter to the current user's documents
results = vectorstore.similarity_search(
    query,
    filter={"owner": current_user_team}
)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

import time
import logging

logger = logging.getLogger("rag")

def ask_with_logging(question: str) -> dict:
    start = time.time()
    result = ask(question)
    elapsed = time.time() - start

    logger.info({
        "query": question,
        "num_chunks": len(result["sources"]),
        "top_score": result["sources"][0][1] if result["sources"] else 0,
        "latency_ms": int(elapsed * 1000),
        "answer_length": len(result["answer"]),
    })
    return result