Fix metadata association to display correct page numbers in sources

app.py

@@ -28,18 +28,20 @@ Question:
 user_role_prompt = UserRolePrompt(user_prompt_template)
 
 class RetrievalAugmentedQAPipeline:
-    def __init__(self, llm: ChatOpenAI(), vector_db_retriever: VectorDatabase, metadata: List[Dict[str, Any]] = None) -> None:
+    def __init__(self, llm: ChatOpenAI(), vector_db_retriever: VectorDatabase, metadata: List[Dict[str, Any]] = None, texts: List[str] = None) -> None:
         self.llm = llm
         self.vector_db_retriever = vector_db_retriever
         self.metadata = metadata or []
         self.text_to_metadata = {}
         
-        # Create lookup for text to metadata
-        if metadata:
-            texts = [key for key in self.vector_db_retriever.vectors.keys()]
+        # Ensure we have the original texts that match the metadata
+        if metadata and texts and len(texts) == len(metadata):
+            # Create a direct mapping from text to its metadata using the original texts
             for i, text in enumerate(texts):
-                if i < len(metadata):
-                    self.text_to_metadata[text] = metadata[i]
+                self.text_to_metadata[text] = metadata[i]
+            print(f"Successfully mapped {len(self.text_to_metadata)} text chunks to metadata")
+        else:
+            print(f"Warning: Metadata mapping not created. Metadata: {len(metadata) if metadata else 0}, Texts: {len(texts) if texts else 0}")
 
     async def arun_pipeline(self, user_query: str):
         context_list = self.vector_db_retriever.search_by_text(user_query, k=4)
@@ -55,7 +57,18 @@ class RetrievalAugmentedQAPipeline:
             if text in self.text_to_metadata:
                 sources.append(self.text_to_metadata[text])
             else:
-                sources.append({"filename": "unknown", "page": "unknown"})
+                # If exact text not found, try finding most similar text
+                # This is a fallback mechanism
+                found = False
+                for orig_text, meta in self.text_to_metadata.items():
+                    # Simple overlap check - if 80% of the text matches
+                    if len(set(text.split()).intersection(set(orig_text.split()))) / max(len(set(text.split())), 1) > 0.8:
+                        sources.append(meta)
+                        found = True
+                        break
+                
+                if not found:
+                    sources.append({"filename": "unknown", "page": "unknown"})
 
         formatted_system_prompt = system_role_prompt.create_message()
 
@@ -85,10 +98,11 @@ def load_preprocessed_data():
     for key, vector in data['vectors'].items():
         vector_db.insert(key, vector)
     
-    # Get metadata if available
+    # Get metadata and original texts if available
     metadata = data.get('metadata', [])
+    texts = data.get('texts', [])
     
-    return vector_db, metadata
+    return vector_db, metadata, texts
 
 @cl.on_chat_start
 async def on_chat_start():
@@ -121,7 +135,7 @@ The application requires preprocessing of PDF documents to build a knowledge bas
 
         # Load pre-processed data
         start_time = time.time()
-        vector_db, metadata = load_preprocessed_data()
+        vector_db, metadata, texts = load_preprocessed_data()
         load_time = time.time() - start_time
         print(f"Loaded vector database in {load_time:.2f} seconds")
         
@@ -131,7 +145,8 @@ The application requires preprocessing of PDF documents to build a knowledge bas
         retrieval_augmented_qa_pipeline = RetrievalAugmentedQAPipeline(
             vector_db_retriever=vector_db,
             llm=chat_openai,
-            metadata=metadata
+            metadata=metadata,
+            texts=texts
         )
         
         # Let the user know that the system is ready