RAG Pipeline: PDF to Searchable Vector Store

Problem Statement

Unstructured documentation scattered across hundreds of PDF manuals, policy documents, architecture decision records, and post-mortems creates significant friction for engineering teams. Traditional keyword search (like grep or standard PDF search tools) frequently fails because it cannot grasp semantic intent, synonyms, or conceptual context.

This script implements an end-to-end Retrieval-Augmented Generation (RAG) pipeline that ingests a directory of PDF documents, parses their textual content, converts text chunks into vector embeddings, and stores them in a persistent local vector database. Once indexed, users can query the vector store in natural language to retrieve relevant document passages and automatically synthesize grounded, contextual answers using OpenAI's language models.

How It Works

The script operates in two main phases: Ingestion and Retrieval/Generation.

1. Ingestion Phase:

   • Parsing: Traverses the target directory and parses PDF files page-by-page using pypdf.
   • Chunking: Applies a sliding-window text chunker with configurable token/character sizes and overlap to maintain context across chunk boundaries.
   • Hashing: Generates deterministic SHA-256 content hashes for each chunk to ensure idempotency and prevent duplicate entries in the database.
   • Embedding & Storage: Batches text chunks, sends them to OpenAI's text-embedding-3-small API, and upserts vectors alongside source metadata (filename, page number, chunk index) into a local ChromaDB instance.

2. Retrieval & Generation Phase:

   • Converts user queries into vector embeddings and performs nearest-neighbor vector search in ChromaDB using cosine similarity.
   • Retrieves the top-$k$ context passages and feeds them alongside the user query to gpt-4o-mini to synthesize an accurate answer with source citations.

[ PDF Directory ] ──> ( pypdf Extractor ) ──> [ Raw Text ] ──> ( Chunking & Hashing )
                                                                      │
                                                                      ▼
[ OpenAI LLM Response ] <── ( gpt-4o-mini ) <── [ Top Contexts ] <── ( ChromaDB Store ) <── ( Embedding API )

Full Script

#!/usr/bin/env python3
"""
RAG Pipeline: PDF to Searchable Vector Store
Parse PDFs, chunk text, generate embeddings via OpenAI, and index into ChromaDB.
Provides both vector similarity search and grounded RAG answer generation.
"""

import argparse
import hashlib
import logging
import os
import sys
from pathlib import Path
from typing import Any, Dict, List, Optional

import chromadb
from chromadb.config import Settings
import openai
from pypdf import PdfReader

# ------------------------------------------------------------------------------
# Configuration & Setup
# ------------------------------------------------------------------------------
DEFAULT_EMBEDDING_MODEL = os.environ.get("EMBEDDING_MODEL", "text-embedding-3-small")
DEFAULT_LLM_MODEL = os.environ.get("LLM_MODEL", "gpt-4o-mini")
DEFAULT_DB_PATH = os.environ.get("CHROMA_DB_PATH", "./chroma_db")
DEFAULT_COLLECTION_NAME = os.environ.get("COLLECTION_NAME", "pdf_knowledge_base")
BATCH_SIZE = 100  # Number of text chunks to embed per API request

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[logging.StreamHandler(sys.stdout)],
)
logger = logging.getLogger("rag_pipeline")


# ------------------------------------------------------------------------------
# Core Components
# ------------------------------------------------------------------------------
class PDFProcessor:
    """Handles text extraction from PDF files."""

    @staticmethod
    def extract_text_from_pdf(file_path: Path) -> List[Dict[str, Any]]:
        """Extract text page by page from a PDF file."""
        pages_content = []
        try:
            reader = PdfReader(str(file_path))
            for page_num, page in enumerate(reader.pages, start=1):
                text = page.extract_text() or ""
                cleaned_text = " ".join(text.split())
                if cleaned_text:
                    pages_content.append(
                        {
                            "page": page_num,
                            "text": cleaned_text,
                            "source": file_path.name,
                        }
                    )
        except Exception as e:
            logger.error(f"Failed to process PDF {file_path}: {e}")
        return pages_content


class TextChunker:
    """Splits raw text into manageable, overlapping chunks."""

    def __init__(self, chunk_size: int = 1000, chunk_overlap: int = 200):
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap

    def split_pages(
        self, pages_data: List[Dict[str, Any]]
    ) -> List[Dict[str, Any]]:
        """Split page contents into smaller chunks with metadata."""
        chunks = []
        for page_info in pages_data:
            text = page_info["text"]
            source = page_info["source"]
            page = page_info["page"]

            start = 0
            text_len = len(text)
            chunk_idx = 0

            while start < text_len:
                end = start + self.chunk_size
                chunk_text = text[start:end]

                # Compute deterministic hash for deduplication ID
                hash_input = f"{source}_p{page}_c{chunk_idx}_{chunk_text}"
                chunk_id = hashlib.sha256(hash_input.encode("utf-8")).hexdigest()

                chunks.append(
                    {
                        "id": chunk_id,
                        "text": chunk_text,
                        "metadata": {
                            "source": source,
                            "page": page,
                            "chunk_index": chunk_idx,
                        },
                    }
                )

                if end >= text_len:
                    break
                start += self.chunk_size - self.chunk_overlap
                chunk_idx += 1

        return chunks


class VectorStoreManager:
    """Manages embedding generation and ChromaDB collection indexing."""

    def __init__(
        self,
        db_path: str,
        collection_name: str,
        embedding_model: str,
        openai_client: openai.OpenAI,
    ):
        self.client = chromadb.PersistentClient(path=db_path)
        self.collection = self.client.get_or_create_collection(
            name=collection_name, metadata={"hnsw:space": "cosine"}
        )
        self.embedding_model = embedding_model
        self.openai_client = openai_client

    def _get_embeddings(self, texts: List[str]) -> List[List[float]]:
        """Fetch vector embeddings from OpenAI API."""
        response = self.openai_client.embeddings.create(
            model=self.embedding_model, input=texts
        )
        return [item.embedding for item in response.data]

    def index_chunks(self, chunks: List[Dict[str, Any]]) -> None:
        """Upsert text chunks into ChromaDB in batches."""
        if not chunks:
            logger.info("No new chunks to index.")
            return

        total_chunks = len(chunks)
        logger.info(f"Indexing {total_chunks} chunks into ChromaDB...")

        for i in range(0, total_chunks, BATCH_SIZE):
            batch = chunks[i : i + BATCH_SIZE]
            ids = [c["id"] for c in batch]
            documents = [c["text"] for c in batch]
            metadatas = [c["metadata"] for c in batch]

            embeddings = self._get_embeddings(documents)

            self.collection.upsert(
                ids=ids,
                embeddings=embeddings,
                documents=documents,
                metadatas=metadatas,
            )
            logger.info(
                f"Indexed batch {i // BATCH_SIZE + 1}/{(total_chunks + BATCH_SIZE - 1) // BATCH_SIZE}"
            )

    def query(self, query_text: str, top_k: int = 3) -> Dict[str, Any]:
        """Query vector database for similar context documents."""
        query_embedding = self._get_embeddings([query_text])[0]
        results = self.collection.query(
            query_embeddings=[query_embedding],
            n_results=top_k,
            include=["documents", "metadatas", "distances"],
        )
        return results


class RAGPipeline:
    """Orchestrates query resolution and prompt generation using OpenAI."""

    def __init__(
        self,
        vector_store: VectorStoreManager,
        openai_client: openai.OpenAI,
        llm_model: str,
    ):
        self.vector_store = vector_store
        self.openai_client = openai_client
        self.llm_model = llm_model

    def answer_question(self, question: str, top_k: int = 3) -> str:
        """Retrieve relevant context and generate an answer using LLM."""
        search_results = self.vector_store.query(question, top_k=top_k)

        documents = search_results.get("documents", [[]])[0]
        metadatas = search_results.get("metadatas", [[]])[0]

        if not documents:
            return "No relevant context found in the vector store."

        context_blocks = []
        for doc, meta in zip(documents, metadatas):
            source_info = f"[Source: {meta['source']}, Page {meta['page']}]"
            context_blocks.append(f"{source_info}\n{doc}")

        context_str = "\n\n---\n\n".join(context_blocks)

        system_prompt = (
            "You are a helpful assistant answering user questions based on the provided context.\n"
            "Only use the information supplied in the context below. If you cannot answer based on\n"
            "the context, state clearly that the context does not contain the necessary information.\n"
            "Always reference source filenames and page numbers in your answer."
        )

        user_prompt = f"Context:\n{context_str}\n\nQuestion: {question}"

        response = self.openai_client.chat.completions.create(
            model=self.llm_model,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ],
            temperature=0.2,
        )

        return response.choices[0].message.content or ""


# ------------------------------------------------------------------------------
# CLI Main Entrypoint
# ------------------------------------------------------------------------------
def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Ingest PDFs into ChromaDB and execute RAG search queries."
    )
    parser.add_argument(
        "--docs",
        type=str,
        help="Path to directory containing PDF files to ingest.",
    )
    parser.add_argument(
        "--db-path",
        type=str,
        default=DEFAULT_DB_PATH,
        help="Directory path for ChromaDB persistence storage.",
    )
    parser.add_argument(
        "--collection",
        type=str,
        default=DEFAULT_COLLECTION_NAME,
        help="Name of the ChromaDB collection.",
    )
    parser.add_argument(
        "--query",
        type=str,
        help="Search query or question to execute against the vector store.",
    )
    parser.add_argument(
        "--chunk-size",
        type=int,
        default=1000,
        help="Max character size per text chunk.",
    )
    parser.add_argument(
        "--chunk-overlap",
        type=int,
        default=200,
        help="Overlap characters between sequential chunks.",
    )
    parser.add_argument(
        "--top-k",
        type=int,
        default=3,
        help="Number of context passages to retrieve.",
    )
    return parser.parse_args()


def main():
    args = parse_args()

    api_key = os.environ.get("OPENAI_API_KEY")
    if not api_key:
        logger.error("OPENAI_API_KEY environment variable is not set.")
        sys.exit(1)

    openai_client = openai.OpenAI(api_key=api_key)

    vector_store = VectorStoreManager(
        db_path=args.db_path,
        collection_name=args.collection,
        embedding_model=DEFAULT_EMBEDDING_MODEL,
        openai_client=openai_client,
    )

    # Ingestion Workflow
    if args.docs:
        docs_path = Path(args.docs)
        if not docs_path.exists() or not docs_path.is_dir():
            logger.error(f"Provided --docs directory does not exist: {args.docs}")
            sys.exit(1)

        pdf_files = list(docs_path.glob("*.pdf"))
        if not pdf_files:
            logger.warning(f"No PDF files found in {args.docs}")
        else:
            logger.info(f"Found {len(pdf_files)} PDF document(s) to process.")
            chunker = TextChunker(
                chunk_size=args.chunk_size, chunk_overlap=args.chunk_overlap
            )
            all_chunks = []

            for pdf_file in pdf_files:
                logger.info(f"Extracting text from: {pdf_file.name}")
                pages_data = PDFProcessor.extract_text_from_pdf(pdf_file)
                chunks = chunker.split_pages(pages_data)
                all_chunks.extend(chunks)

            vector_store.index_chunks(all_chunks)
            logger.info("Ingestion phase completed successfully.")

    # Query / Generation Workflow
    if args.query:
        logger.info(f"Executing query: '{args.query}'")
        pipeline = RAGPipeline(
            vector_store=vector_store,
            openai_client=openai_client,
            llm_model=DEFAULT_LLM_MODEL,
        )
        answer = pipeline.answer_question(args.query, top_k=args.top_k)
        print("\n" + "=" * 60)
        print("RAG ANSWER:")
        print("=" * 60)
        print(answer)
        print("=" * 60 + "\n")
    elif not args.docs:
        logger.warning(
            "No action specified. Pass --docs to ingest files or --query to search."
        )


if __name__ == "__main__":
    main()

Configuration

Set the following environment variables before executing the pipeline:

Usage Examples

1. Ingesting PDFs into Vector Store

Ingest all PDFs located in ./documents into local storage:

export OPENAI_API_KEY="sk-proj-..."
python3 rag_pipeline.py --docs ./documents/

Output:

2026-03-30 10:15:00 [INFO] Found 2 PDF document(s) to process.
2026-03-30 10:15:00 [INFO] Extracting text from: kubernetes_guide.pdf
2026-03-30 10:15:01 [INFO] Extracting text from: incident_postmortem.pdf
2026-03-30 10:15:02 [INFO] Indexing 42 chunks into ChromaDB...
2026-03-30 10:15:03 [INFO] Indexed batch 1/1
2026-03-30 10:15:03 [INFO] Ingestion phase completed successfully.

2. Querying the Index (Ingest + Query in One Pass)

Ingest new PDFs and answer a question immediately:

python3 rag_pipeline.py \
  --docs ./documents/ \
  --query "What was the root cause of the outage mentioned in the postmortem?" \
  --top-k 3

Output:

... [Ingestion Logs] ...
2026-03-30 10:16:10 [INFO] Executing query: 'What was the root cause of the outage mentioned in the postmortem?'

============================================================
RAG ANSWER: