Update src/RAG_builder.py

src/RAG_builder.py

@@ -34,16 +34,40 @@ class ProjectRAGGraph:
         self.memory = MemorySaver()
         self.workflow = self._build_graph()
 
-    def process_documents(self, pdf_paths):
-        self.pdf_count = len(pdf_paths) # Track how many PDFs were uploaded
+    def process_documents(self, pdf_paths, original_names=None):
+        self.pdf_count = len(pdf_paths)
         all_docs = []
-        for path in pdf_paths:
+        
+        # Iterate through paths and original names simultaneously
+        for i, path in enumerate(pdf_paths):
             loader = PyPDFLoader(path)
-            all_docs.extend(loader.load())
+            docs = loader.load()
+            
+            # If original names are provided, overwrite the 'source' metadata
+            if original_names and i < len(original_names):
+                for doc in docs:
+                    doc.metadata["source"] = original_names[i]
+            
+            all_docs.extend(docs)
+                
+        # Split documents after metadata has been corrected
+        splits = RecursiveCharacterTextSplitter(
+            chunk_size=500, 
+            chunk_overlap=100
+        ).split_documents(all_docs)
         
-        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):
+    #     self.pdf_count = len(pdf_paths) # Track how many PDFs were uploaded
+    #     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 ---")