Mehramd/multi-file-rag-chat

1from langchain_text_splitters import RecursiveCharacterTextSplitter
2from langchain_community.vectorstores import FAISS
3from langchain_huggingface import HuggingFaceEmbeddings
4from langchain_groq import ChatGroq
5from langchain_core.prompts import ChatPromptTemplate
6from langchain_core.output_parsers import StrOutputParser
7from pathlib import Path
8from config import GROQ_API_KEY, GROQ_MODEL
9from config import EMBEDDING_MODEL
10
11embeddings = HuggingFaceEmbeddings(
12    model_name=EMBEDDING_MODEL
13)
14
15def call_llm():
16    return ChatGroq(
17        model=GROQ_MODEL,
18        api_key=GROQ_API_KEY,
19        temperature=0.1,
20    )
21
22VECTOR_DB_PATH = Path("vector_db")
23
24def build_vector_store(documents):
25  splitter = RecursiveCharacterTextSplitter(
26      chunk_size=500,
27      chunk_overlap=100
28  )
29  chunks = splitter.split_documents(documents)
30  vector_store = FAISS.from_documents(chunks,embeddings)
31  vector_store.save_local(str(VECTOR_DB_PATH))
32  return vector_store
33
34def ask_question(question, vector_store):
35  retrieved_docs = vector_store.similarity_search(question, k=2)
36  context = "\n\n".join([doc.page_content for doc in retrieved_docs])
37
38  prompt = ChatPromptTemplate.from_messages(
39    [
40        ("system",
41         """
42         You are a helpful AI assistant.
43
44            Answer ONLY from the provided context.
45
46            If the answer is not present, say:
47            "I don't know from the uploaded PDF.
48            """
49         ),
50        ("user",
51         """
52
53            Context:
54            {context}
55
56            Question:
57            {question}
58            """
59
60         )
61    ]
62)
63
64  llm = call_llm()
65
66  chain = prompt | llm | StrOutputParser()
67  answer = chain.invoke(
68      {
69          "context": context,
70          "question": question
71      }
72  )
73
74  sources = []
75  for doc in retrieved_docs:
76        source = doc.metadata.get("source", "unknown")
77        page = doc.metadata.get("page")
78        label = f"{source}, page {page}" if page else source
79        if label not in sources:
80            sources.append(label)
81
82  return answer, sources
83