Update src/RAG_builder.py

src/RAG_builder.py

@@ -29,46 +29,28 @@ class ProjectRAGGraph:
             api_key="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" # Keep your API keys safe!
         )
         self.vector_store = None
-        
+        self.pdf_count = 0
         # 2. Initialize Memory Checkpointer
         self.memory = MemorySaver()
         self.workflow = self._build_graph()
 
-    def process_documents(self, pdf_paths_with_names: list[tuple[str, str]]):
-        """
-        Expects a list of tuples: [(temp_path, original_name), ...]
-        """
+    def process_documents(self, pdf_paths):
+        self.pdf_count = len(pdf_paths) # Track how many PDFs were uploaded
         all_docs = []
-        for temp_path, original_name in pdf_paths_with_names:
-            loader = PyPDFLoader(temp_path)
-            docs = loader.load()
-            
-            # Override the metadata source with the original filename
-            for doc in docs:
-                doc.metadata["source"] = original_name
-                
-            all_docs.extend(docs)
-        
-        splits = RecursiveCharacterTextSplitter(
-            chunk_size=500, 
-            chunk_overlap=100
-        ).split_documents(all_docs)
+        for path in pdf_paths:
+            loader = PyPDFLoader(path)
+            all_docs.extend(loader.load())
         
+        splits = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100).split_documents(all_docs)
         self.vector_store = FAISS.from_documents(splits, self.embeddings)
-    
-    # def process_documents(self, pdf_paths):
-    #     all_docs = []
-    #     for path in pdf_paths:
-    #         loader = PyPDFLoader(path)
-    #         all_docs.extend(loader.load())
-        
-    #     splits = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100).split_documents(all_docs)
-    #     self.vector_store = FAISS.from_documents(splits, self.embeddings)
 
     # --- GRAPH NODES ---
     def retrieve(self, state: GraphState):
         print("--- RETRIEVING ---")
-        retriever = self.vector_store.as_retriever(search_type="mmr", search_kwargs={"k": 5, "lambda_mult":0.25})
+        # Calculate dynamic k
+        dynamic_k = self.pdf_count + 2
+        k_value = max(1, dynamic_k)
+        retriever = self.vector_store.as_retriever(search_type="mmr", search_kwargs={"k": k_value, "lambda_mult":0.25})
         documents = retriever.invoke(state["question"])
         return {"context": documents}