RAG: Chat with your PDF

prompt = ChatPromptTemplate.from_messages([
    ("system", "You are a helpful assistant. Answer ONLY from the context below. "
               "If the answer is not in the context, say: I don't know based on the document.\n\n"
               "Context:\n{context}"),
    ("human", "{question}"),
])

rag_chain = (
    {"context": retriever | format_docs, "question": RunnablePassthrough()}   # 1 fetch context, keep question
    | prompt                                                                  # 2 fill the template
    | llm                                                                     # 3 answer from context
    | StrOutputParser()                                                       # 4 plain text
)

import warnings; warnings.filterwarnings("ignore")
from standins import KeywordEmbeddings, RuleModel
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough, RunnableParallel

pages = PyPDFLoader("novabank_handbook.pdf").load()                                         # 1 Load
chunks = RecursiveCharacterTextSplitter(chunk_size=200, chunk_overlap=30).split_documents(pages)   # 2 Split
vector_store = FAISS.from_documents(chunks, KeywordEmbeddings())                             # 3 Embed + 4 Store
retriever = vector_store.as_retriever(search_kwargs={"k": 2})                                # 5 Retrieve
print(f"indexed: {len(pages)} pages -> {len(chunks)} chunks")

class ContextOnly(RuleModel):
    """Stand-in LLM that obeys 'answer only from context': it quotes the best context line or says I don't know."""
    def _pick(self, messages):
        import re
        sys_txt, q = messages[0].content, messages[-1].content.lower()
        words = {w for w in re.findall(r"[a-z0-9]+", q) if len(w) > 2} - {"what", "does", "when", "who", "the", "for", "how", "mean", "available"}
        best = max((l for l in sys_txt.splitlines() if l and not l.startswith("You")), key=lambda l: len(words & set(re.findall(r"[a-z0-9]+", l.lower()))), default="")
        hits = len(words & set(re.findall(r"[a-z0-9]+", best.lower())))
        return best.split("] ", 1)[-1] if hits >= 1 else "I don't know based on the document."
llm = ContextOnly()

prompt = ChatPromptTemplate.from_messages([
    ("system", "You are a helpful assistant. Answer ONLY from the context below. "
               "If the answer is not in the context, say: I don't know based on the document.\n\nContext:\n{context}"),
    ("human", "{question}")])
def format_docs(docs): return "\n".join(f"[page {d.metadata['page']}] {l}" for d in docs for l in d.page_content.splitlines()[1:] or [d.page_content])

rag_chain = {"context": retriever | format_docs, "question": RunnablePassthrough()} | prompt | llm | StrOutputParser()  # 6 Generate
for q in ["What is the UPI limit for a new payee?", "What does error U30 mean?", "Who won the cricket world cup 2011?"]:
    print("\nQ:", q, "\nA:", rag_chain.invoke(q))

rag_with_sources = RunnableParallel({"docs": retriever, "question": RunnablePassthrough()}).assign(
    answer=(lambda x: {"context": format_docs(x["docs"]), "question": x["question"]}) | prompt | llm | StrOutputParser())
r = rag_with_sources.invoke("When is video KYC available?")
print("\nA:", r["answer"], "| sources: pages", sorted({d.metadata["page"] + 1 for d in r["docs"]}))