Update app/policy_vector_db.py

app/policy_vector_db.py

@@ -23,7 +23,9 @@ class PolicyVectorDB:
         
         # Using a powerful open-source embedding model.
         # Change 'cpu' to 'cuda' if a GPU is available for significantly faster embedding.
+        logger.info("Loading embedding model 'BAAI/bge-large-en-v1.5'. This may take a moment...")
         self.embedding_model = SentenceTransformer('BAAI/bge-large-en-v1.5', device='cpu') 
+        logger.info("Embedding model loaded successfully.")
         
         self.collection = None # Initialize collection as None for lazy loading
         self.top_k_default = top_k_default
@@ -69,7 +71,7 @@ class PolicyVectorDB:
         logger.info(f"Adding {len(new_chunks)} new chunks to the vector database...")
         
         # Process in batches for efficiency
-        batch_size = 64
+        batch_size = 32 # Reduced batch size for potentially large embeddings
         for i in range(0, len(new_chunks), batch_size):
             batch = new_chunks[i:i + batch_size]
             
@@ -77,7 +79,8 @@ class PolicyVectorDB:
             texts = [chunk['text'] for chunk in batch]
             metadatas = [self._flatten_metadata(chunk.get('metadata', {})) for chunk in batch]
             
-            embeddings = self.embedding_model.encode(texts, show_progress_bar=False).tolist()
+            # For BGE models, it's recommended not to add instructions to the document embeddings
+            embeddings = self.embedding_model.encode(texts, normalize_embeddings=True, show_progress_bar=False).tolist()
             
             collection.add(ids=ids, embeddings=embeddings, documents=texts, metadatas=metadatas)
             logger.info(f"Added batch {i//batch_size + 1}/{(len(new_chunks) + batch_size - 1) // batch_size}")
@@ -90,19 +93,26 @@ class PolicyVectorDB:
         Returns a list of results filtered by a relevance threshold.
         """
         collection = self._get_collection()
-        query_embedding = self.embedding_model.encode([query_text]).tolist()
+        
+        # ✅ IMPROVEMENT: Add the recommended instruction prefix for BGE retrieval models.
+        instructed_query = f"Represent this sentence for searching relevant passages: {query_text}"
+        
+        # ✅ IMPROVEMENT: Normalize embeddings for more accurate similarity search.
+        query_embedding = self.embedding_model.encode([instructed_query], normalize_embeddings=True).tolist()
+        
         k = top_k if top_k is not None else self.top_k_default
         
         # Retrieve more results initially to allow for filtering
         results = collection.query(
             query_embeddings=query_embedding,
-            n_results=k * 2,
+            n_results=k * 2, # Retrieve more to filter by threshold
             include=["documents", "metadatas", "distances"]
         )
         
         search_results = []
         if results and results.get('documents') and results['documents'][0]:
             for i, doc in enumerate(results['documents'][0]):
+                # The distance for normalized embeddings is often interpreted as 1 - cosine_similarity
                 relevance_score = 1 - results['distances'][0][i]
                 
                 if relevance_score >= self.relevance_threshold:
@@ -112,6 +122,7 @@ class PolicyVectorDB:
                         'relevance_score': relevance_score
                     })
         
+        # Sort by relevance score and return the top_k results
         return sorted(search_results, key=lambda x: x['relevance_score'], reverse=True)[:k]
 
 def ensure_db_populated(db_instance: PolicyVectorDB, chunks_file_path: str) -> bool:
@@ -145,4 +156,4 @@ def ensure_db_populated(db_instance: PolicyVectorDB, chunks_file_path: str) -> b
         return True
     except Exception as e:
         logger.error(f"An error occurred during DB population check: {e}", exc_info=True)
-        return False
+        return False