pdf_rag.py

import fitz
from langchain_text_splitters import RecursiveCharacterTextSplitter
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct
from sentence_transformers import SentenceTransformer

def extract_text_from_pdf(pdf_path:str) -> str:
    print(f"Extracting text from {pdf_path}")
    doc = fitz.open(pdf_path)
    full_text = ""
    for page_num, page in enumerate(doc):
        full_text += page.get_text("text") + "\n"
    return full_text

def chunk_text(text:str) -> list[str]:
    splitter = RecursiveCharacterTextSplitter(
        chunk_size = 500,
        chunk_overlap = 50,
        separators=["\n\n", "\n", ".", " ", ""]
    )
    return splitter.split_text(text)

def run_mini_rag(pdf_path:str, user_query:str):
    client = QdrantClient(path="./qdrant_pdf_db")
    model = SentenceTransformer("all-MiniLM-l6-v2")
    COLLECTION_NAME = "pdf_collection"

    if client.collection_exists(COLLECTION_NAME):
        client.delete_collection(COLLECTION_NAME)
    
    client.create_collection(
        collection_name=COLLECTION_NAME,
        vectors_config=VectorParams(size= 384, distance=Distance.COSINE)
    )

    raw_text = extract_text_from_pdf(pdf_path)
    chunks = chunk_text(raw_text)

    print(f"Embedding {len(chunks)} chunks into vector DB")
    points = []
    for i, chunk in enumerate(chunks):
        vector = model.encode(chunk).tolist()
        points.append(PointStruct(
            id = i + 1,
            vector=vector,
            payload={"text": chunk, "page_estimate": i}
        ))

    client.upsert(collection_name=COLLECTION_NAME, points=points)
    print("\n Searching DB for : {user_query}")
    query_vector = model.encode(user_query).tolist()
    search_results = client.query_points(
        collection_name=COLLECTION_NAME,
        query = query_vector,
        limit=2
    )
    retrieved_context = "\n---\n".join([hit.payload["text"] for hit in search_results.points])
    final_prompt = f"""
    You are an expert assistant. Answer the user's question using ONLY the provided context.
    If the answer is not in the context, say "I don't have enough verified information to answer this."
    
    CONTEXT:
    {retrieved_context}
    
    USER QUESTION: {user_query}
    """
    
    print("\n=======================================================")
    print("FINAL PROMPT READY TO BE SENT TO AN LLM (like Llama-3):")
    print(final_prompt)

if __name__ == "__main__":
    pdf_file = "sample.pdf"
    question = "What are the Key Takeaway from this document?"

    run_mini_rag(pdf_file, question)