5. RAG and Q&A

Description

The RAGModule is the component that gives the Document Analyzer memory. While the other modules process a document once and produce immediate results, RAG indexes the content so you can ask questions later — even across multiple documents. A user uploads an invoice today, a contract tomorrow, and can ask "how much did we pay Tech Solutions in total?" and get an answer with sources.

Why it matters: Without RAG, each question requires reprocessing the entire document. With RAG, the document is processed once, indexed in a vector database, and questions are answered in milliseconds by querying the most relevant chunks. This is especially valuable when you have dozens or hundreds of indexed documents.

Connection with the module: In Module 6 you built a basic RAG system with ChromaDB and embeddings. Here you integrate it into the Document Analyzer with improvements: indexing of text AND image descriptions (multimodal RAG), metadata per page and document, hybrid retrieval with filters, and answer generation with cited sources.

RAGModule Architecture

Responsibilities

RAGModule
├── index()          → Index a processed document in ChromaDB
├── query()          → Search relevant chunks and generate an answer
├── delete()         → Remove a document from the index
├── list_documents() → List indexed documents
└── get_stats()      → Index statistics

Indexing flow

ProcessedDocument
    │
    ├── Text pages → TextChunker → Chunks with metadata
    │                                 │
    │                                 ▼
    │                        ChromaDB.add(documents, metadatas, ids)
    │
    └── Image pages → VisionAnalyzer.describe_for_rag()
                         │
                         ▼
                    Text descriptions
                         │
                         ▼
               ChromaDB.add(descriptions, metadatas, ids)

Q&A flow

User's question
    │
    ▼
ChromaDB.query(question, n_results=5)
    │
    ▼
Relevant chunks + metadata (page, type, document)
    │
    ▼
LLM: "Given this context, answer the question"
    │
    ▼
Answer + cited sources

Data Model for RAG

Indexed chunk

from pydantic import BaseModel
from typing import Optional


class IndexedChunk(BaseModel):
    chunk_id: str
    doc_id: str
    text: str
    page_number: int
    content_type: str  # "text" | "image_description"
    metadata: dict = {}


class QAResult(BaseModel):
    question: str
    answer: str
    sources: list[dict] = []
    chunks_used: int = 0
    confidence: Optional[float] = None

Metadata structure in ChromaDB

Each chunk is stored with metadata that lets you filter by document, page and type:

{
    "doc_id": "a1b2c3d4",
    "page_number": 3,
    "content_type": "text",          # "text" | "image_description"
    "filename": "invoice_march.pdf",
    "document_type": "invoice",       # VisionAnalyzer classification
    "chunk_index": 2,
    "char_count": 1423
}

Implementation: RAGModule

Complete class

import logging
import os
from typing import Optional

import chromadb
from chromadb.utils import embedding_functions
from openai import OpenAI

logger = logging.getLogger(__name__)

DEFAULT_COLLECTION = "document_analyzer"
DEFAULT_N_RESULTS = 5
CHUNK_SIZE = 1500
CHUNK_OVERLAP = 200
MIN_CHUNK_SIZE = 100


class RAGModule:
    def __init__(
        self,
        collection_name: str = DEFAULT_COLLECTION,
        persist_directory: Optional[str] = None
    ):
        self.openai_client = OpenAI()

        self.embedding_fn = embedding_functions.OpenAIEmbeddingFunction(
            api_key=os.getenv("OPENAI_API_KEY"),
            model_name="text-embedding-3-small"
        )

        if persist_directory:
            self.chroma_client = chromadb.PersistentClient(path=persist_directory)
        else:
            self.chroma_client = chromadb.Client()

        self.collection = self.chroma_client.get_or_create_collection(
            name=collection_name,
            embedding_function=self.embedding_fn,
            metadata={"hnsw:space": "cosine"}
        )

        logger.info(
            f"RAGModule initialized. Collection: {collection_name}, "
            f"existing documents: {self.collection.count()}"
        )

    def index(
        self,
        doc_id: str,
        content,
        filename: str = "",
        document_type: str = "other",
        image_descriptions: Optional[list[str]] = None
    ) -> dict:
        chunks_indexed = 0

        if content.has_text_pages and content.full_text:
            text_chunks = self._chunk_text(content.full_text, doc_id)
            if text_chunks:
                self.collection.add(
                    documents=[c["text"] for c in text_chunks],
                    ids=[c["id"] for c in text_chunks],
                    metadatas=[{
                        "doc_id": doc_id,
                        "page_number": 0,
                        "content_type": "text",
                        "filename": filename,
                        "document_type": document_type,
                        "chunk_index": c["index"],
                        "char_count": len(c["text"])
                    } for c in text_chunks]
                )
                chunks_indexed += len(text_chunks)
                logger.info(f"Indexed {len(text_chunks)} text chunks for {doc_id}")

        if image_descriptions:
            desc_ids = [f"{doc_id}_imgdesc_{i}" for i in range(len(image_descriptions))]
            self.collection.add(
                documents=image_descriptions,
                ids=desc_ids,
                metadatas=[{
                    "doc_id": doc_id,
                    "page_number": i + 1,
                    "content_type": "image_description",
                    "filename": filename,
                    "document_type": document_type,
                    "chunk_index": i,
                    "char_count": len(desc)
                } for i, desc in enumerate(image_descriptions)]
            )
            chunks_indexed += len(image_descriptions)
            logger.info(f"Indexed {len(image_descriptions)} image descriptions for {doc_id}")

        return {
            "doc_id": doc_id,
            "chunks_indexed": chunks_indexed,
            "total_in_collection": self.collection.count()
        }

    def query(
        self,
        question: str,
        doc_id: Optional[str] = None,
        n_results: int = DEFAULT_N_RESULTS,
        content_type: Optional[str] = None
    ) -> QAResult:
        where_filter = self._build_filter(doc_id, content_type)

        query_kwargs = {
            "query_texts": [question],
            "n_results": min(n_results, self.collection.count() or 1),
        }
        if where_filter:
            query_kwargs["where"] = where_filter

        try:
            results = self.collection.query(**query_kwargs)
        except Exception as e:
            logger.error(f"Error in ChromaDB query: {e}")
            return QAResult(
                question=question,
                answer="Error searching the index.",
                sources=[], chunks_used=0
            )

        if not results["documents"] or not results["documents"][0]:
            return QAResult(
                question=question,
                answer="No relevant documents found to answer the question.",
                sources=[], chunks_used=0
            )

        documents = results["documents"][0]
        metadatas = results["metadatas"][0] if results.get("metadatas") else [{}] * len(documents)
        distances = results["distances"][0] if results.get("distances") else [0] * len(documents)

        context = self._build_context(documents, metadatas)
        answer = self._generate_answer(question, context)

        sources = []
        for i, (doc_text, meta, dist) in enumerate(zip(documents, metadatas, distances)):
            sources.append({
                "text": doc_text[:200] + "..." if len(doc_text) > 200 else doc_text,
                "page": meta.get("page_number", "?"),
                "type": meta.get("content_type", "?"),
                "filename": meta.get("filename", "?"),
                "relevance": round(1 - dist, 3) if dist else None
            })

        avg_relevance = sum(s["relevance"] for s in sources if s["relevance"]) / len(sources) if sources else 0

        return QAResult(
            question=question,
            answer=answer,
            sources=sources,
            chunks_used=len(documents),
            confidence=round(avg_relevance, 2)
        )

    def delete(self, doc_id: str) -> dict:
        all_items = self.collection.get(where={"doc_id": doc_id})
        if all_items["ids"]:
            self.collection.delete(ids=all_items["ids"])
            logger.info(f"Deleted {len(all_items['ids'])} chunks from {doc_id}")
            return {"deleted": len(all_items["ids"]), "doc_id": doc_id}
        return {"deleted": 0, "doc_id": doc_id}

    def list_documents(self) -> list[dict]:
        all_items = self.collection.get()
        doc_ids = set()
        docs = {}

        for meta in all_items.get("metadatas", []):
            did = meta.get("doc_id", "unknown")
            if did not in docs:
                docs[did] = {
                    "doc_id": did,
                    "filename": meta.get("filename", "?"),
                    "document_type": meta.get("document_type", "?"),
                    "chunks": 0
                }
            docs[did]["chunks"] += 1

        return list(docs.values())

    def get_stats(self) -> dict:
        total = self.collection.count()
        docs = self.list_documents()
        return {
            "total_chunks": total,
            "total_documents": len(docs),
            "documents": docs
        }

    def _chunk_text(self, text: str, doc_id: str) -> list[dict]:
        chunks = []
        paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()]
        current = ""
        idx = 0

        for para in paragraphs:
            if len(current) + len(para) + 2 <= CHUNK_SIZE:
                current += ("\n\n" + para if current else para)
            else:
                if len(current) >= MIN_CHUNK_SIZE:
                    chunks.append({
                        "id": f"{doc_id}_chunk_{idx}",
                        "text": current,
                        "index": idx
                    })
                    idx += 1
                overlap = current[-CHUNK_OVERLAP:] if current else ""
                current = overlap + "\n\n" + para if overlap else para

        if len(current) >= MIN_CHUNK_SIZE:
            chunks.append({
                "id": f"{doc_id}_chunk_{idx}",
                "text": current,
                "index": idx
            })

        return chunks

    def _build_filter(
        self, doc_id: Optional[str], content_type: Optional[str]
    ) -> Optional[dict]:
        conditions = []
        if doc_id:
            conditions.append({"doc_id": {"$eq": doc_id}})
        if content_type:
            conditions.append({"content_type": {"$eq": content_type}})

        if len(conditions) == 0:
            return None
        if len(conditions) == 1:
            return conditions[0]
        return {"$and": conditions}

    def _build_context(self, documents: list[str], metadatas: list[dict]) -> str:
        parts = []
        for i, (doc, meta) in enumerate(zip(documents, metadatas)):
            source_info = f"[Source {i+1}: {meta.get('filename', '?')}, page {meta.get('page_number', '?')}]"
            parts.append(f"{source_info}\n{doc}")
        return "\n\n---\n\n".join(parts)

    def _generate_answer(self, question: str, context: str) -> str:
        r = self.openai_client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {
                    "role": "system",
                    "content": (
                        "You are an assistant that answers questions about documents. "
                        "Answer ONLY based on the provided context. "
                        "If the information is not in the context, say that you don't have enough information. "
                        "Cite the sources when relevant (e.g., '[Source 1]')."
                    )
                },
                {
                    "role": "user",
                    "content": f"Context:\n{context}\n\nQuestion: {question}"
                }
            ],
            max_tokens=500,
            temperature=0
        )
        return r.choices[0].message.content

Multimodal Indexing

The concept

Multimodal RAG indexes both direct text and generated descriptions of images. When a user asks about a table that was on a scanned page, the system finds the description of that image and uses it as context.

10-page document
├── Pages 1-7 (text) → 12 text chunks indexed
└── Pages 8-10 (scanned) → 3 image descriptions indexed

Total in ChromaDB: 15 chunks

Integration with VisionAnalyzer

def index_full_document(
    rag: RAGModule,
    analyzer: VisionAnalyzer,
    doc_id: str,
    content,
    filename: str,
    document_type: str
) -> dict:
    image_descriptions = None

    if content.has_image_pages:
        images = content.get_images_for_vision()
        image_descriptions = analyzer.describe_for_rag(images)

    return rag.index(
        doc_id=doc_id,
        content=content,
        filename=filename,
        document_type=document_type,
        image_descriptions=image_descriptions
    )

Retrieval with Filters

Filter use cases

Example of filtered queries

rag = RAGModule(persist_directory="./chroma_data")

result = rag.query(
    question="What is the total of the March invoice?",
    doc_id="invoice_march_001"
)
print(f"Answer: {result.answer}")
print(f"Sources: {len(result.sources)}")

result_all = rag.query(
    question="How much have we paid in total this quarter?",
    doc_id=None  # searches across all documents
)
print(f"Answer: {result_all.answer}")

result_images = rag.query(
    question="What do the report charts show?",
    content_type="image_description"
)
print(f"Answer: {result_images.answer}")

Index Persistence

Why persist

Without persistence, the index is lost when the service restarts. Each document would be re-indexed on every request.

# No persistence (development only)
rag = RAGModule()

# With persistence (production)
rag = RAGModule(persist_directory="./chroma_data")

Managing the persistent index

stats = rag.get_stats()
print(f"Total chunks: {stats['total_chunks']}")
print(f"Total documents: {stats['total_documents']}")

for doc in stats["documents"]:
    print(f"  {doc['doc_id']}: {doc['filename']} ({doc['chunks']} chunks)")

rag.delete("old_document_001")

Troubleshooting

"The answers aren't relevant"

Probable cause: Chunks that are too large dilute relevance. Or low-quality embeddings.

Solution: Reduce CHUNK_SIZE:

CHUNK_SIZE = 800      # smaller = more precise
CHUNK_OVERLAP = 100   # adjust overlap

Or increase n_results to give the LLM more context:

result = rag.query(question="...", n_results=10)

"ChromaDB throws a dimension error"

Probable cause: You mixed embeddings from different models in the same collection.

Solution: Create a new collection or delete the existing one:

self.chroma_client.delete_collection("document_analyzer")
self.collection = self.chroma_client.create_collection(
    name="document_analyzer",
    embedding_function=self.embedding_fn
)

"Indexing is slow"

Probable cause: Many small chunks generate many embedding calls.

Solution: ChromaDB sends embeddings in batches automatically, but you can index chunks in batches:

BATCH_SIZE = 100

for i in range(0, len(all_chunks), BATCH_SIZE):
    batch = all_chunks[i:i + BATCH_SIZE]
    self.collection.add(
        documents=[c["text"] for c in batch],
        ids=[c["id"] for c in batch],
        metadatas=[c["meta"] for c in batch]
    )

"It can't find recently indexed documents"

Probable cause: Using in-memory ChromaDB and the process restarted.

Solution: Use PersistentClient:

self.chroma_client = chromadb.PersistentClient(path="./chroma_data")

Using the RAGModule

Complete end-to-end example

processor = DocumentProcessor()
analyzer = VisionAnalyzer()
rag = RAGModule(persist_directory="./chroma_data")

doc = processor.process("service_contract.pdf")
doc_type = analyzer.classify(doc)

image_descriptions = None
if doc.has_image_pages:
    image_descriptions = analyzer.describe_for_rag(doc.get_images_for_vision())

index_result = rag.index(
    doc_id="contract_001",
    content=doc,
    filename="service_contract.pdf",
    document_type=doc_type,
    image_descriptions=image_descriptions
)
print(f"Indexed: {index_result['chunks_indexed']} chunks")

qa = rag.query(
    question="What are the parties to the contract?",
    doc_id="contract_001"
)
print(f"\nQuestion: {qa.question}")
print(f"Answer: {qa.answer}")
print(f"Confidence: {qa.confidence}")
print(f"Sources used: {qa.chunks_used}")
for source in qa.sources:
    print(f"  - [{source['type']}] Page {source['page']}: {source['text'][:80]}...")

Expected output

Indexed: 8 chunks

Question: What are the parties to the contract?
Answer: The parties to the contract are: (1) Tech Solutions S.A. as service
provider, and (2) Company ABC S.A. de C.V. as client [Source 1].
Confidence: 0.87
Sources used: 5
  - [text] Page ?: SERVICE AGREEMENT entered into by...
  - [text] Page ?: The parties agree to the following terms...

Exercises

Exercise 1: Answer with formatted sources

Modify query() so that the answer includes the cited sources in a readable format. The LLM must cite "[Source N]" in its answer, and the result must include a mapping of sources with their text and metadata.

def query_with_formatted_sources(
    self,
    question: str,
    doc_id: Optional[str] = None,
    n_results: int = 5
) -> dict:
    result = self.query(question, doc_id, n_results)

    source_map = {}
    for i, source in enumerate(result.sources):
        key = f"Source {i + 1}"
        source_map[key] = {
            "text": source["text"],
            "file": source["filename"],
            "page": source["page"],
            "type": source["type"],
            "relevance": source.get("relevance")
        }

    return {
        "question": result.question,
        "answer": result.answer,
        "confidence": result.confidence,
        "cited_sources": source_map,
        "total_sources": len(source_map)
    }


rag = RAGModule()
formatted = rag.query_with_formatted_sources(
    "What is the signing date?", doc_id="contract_001"
)
print(f"Answer: {formatted['answer']}")
print(f"\nCited sources:")
for key, src in formatted["cited_sources"].items():
    print(f"  [{key}] {src['file']}, page {src['page']}: {src['text'][:60]}...")

Exercise 2: Cross-document search

Implement a function that searches across all indexed documents and groups the sources by document. Useful for questions like "how much did we pay in total this quarter?" that require information from multiple invoices.

def cross_document_query(
    self,
    question: str,
    n_results: int = 10
) -> dict:
    result = self.query(question, doc_id=None, n_results=n_results)

    by_document: dict[str, list] = {}
    for source in result.sources:
        doc_key = source.get("filename", "unknown")
        if doc_key not in by_document:
            by_document[doc_key] = []
        by_document[doc_key].append(source)

    return {
        "question": result.question,
        "answer": result.answer,
        "confidence": result.confidence,
        "documents_queried": len(by_document),
        "sources_by_document": {
            doc: {
                "source_count": len(sources),
                "sources": [s["text"][:100] for s in sources]
            }
            for doc, sources in by_document.items()
        }
    }


rag = RAGModule(persist_directory="./chroma_data")
cross = rag.cross_document_query("How much did we pay in total this quarter?")
print(f"Answer: {cross['answer']}")
print(f"Documents queried: {cross['documents_queried']}")
for doc_name, info in cross["sources_by_document"].items():
    print(f"  {doc_name}: {info['source_count']} sources")

Summary

• The RAGModule indexes documents for Q&A: text as chunks, images as text descriptions.
• It uses ChromaDB with OpenAI embeddings (text-embedding-3-small).
• Retrieval with filters: by document, by content type, or without a filter (cross-document).
• Answer generation: LLM with retrieved context, instructions to cite sources.
• Persistence: PersistentClient to avoid losing the index between restarts.
• Management: list documents, delete, index statistics.
• Multimodal indexing enables Q&A about visual content through descriptions.

Additional Resources

1. ChromaDB Documentation — Complete reference
2. OpenAI Embeddings — Embedding models
3. RAG Best Practices — RAG patterns with LangChain
4. Module 6 of this guide — Multimodal RAG foundation