Update app/policy_vector_db.py

app/policy_vector_db.py

@@ -5,16 +5,24 @@ from typing import List, Dict
 from sentence_transformers import SentenceTransformer
 import chromadb
 from chromadb.config import Settings
+import logging
+
+logger = logging.getLogger("app") # Use the same logger as app.py
 
 class PolicyVectorDB:
     def __init__(self, persist_directory: str, top_k_default: int = 5, relevance_threshold: float = 0.5):
         self.persist_directory = persist_directory
+        # Allow reset is useful for development, can be removed in production if not needed
         self.client = chromadb.PersistentClient(path=persist_directory, settings=Settings(allow_reset=True))
         self.collection_name = "neepco_dop_policies"
+
+        # IMPORTANT: Keeping BAAI/bge-large-en-v1.5 as per your requirement for accuracy.
+        # Be aware this will be slow on CPU.
         self.embedding_model = SentenceTransformer('BAAI/bge-large-en-v1.5', device='cuda' if torch.cuda.is_available() else 'cpu')
+        
         self.collection = None
         self.top_k_default = top_k_default
-        self.relevance_threshold = relevance_threshold
+        self.relevance_threshold = relevance_threshold # Note: set to 0 in app.py
 
     def _get_collection(self):
         if self.collection is None:
@@ -25,58 +33,100 @@ class PolicyVectorDB:
         return self.collection
 
     def _flatten_metadata(self, metadata: Dict) -> Dict:
+        # ChromaDB requires metadata values to be JSON-serializable strings, ints, floats, or bools.
+        # Ensuring all values are string for consistency and compatibility.
         return {key: str(value) for key, value in metadata.items()}
 
     def add_chunks(self, chunks: List[Dict]):
         collection = self._get_collection()
         if not chunks:
-            print("No chunks provided to add.")
+            logger.info("No chunks provided to add.")
             return
-        existing_ids = set(collection.get()['ids'])
-        new_chunks = [chunk for chunk in chunks if chunk.get('id') not in existing_ids]
+
+        # Fetch existing IDs to avoid adding duplicates
+        existing_ids = set()
+        try:
+            # This can be slow for very large collections, consider optimizing if needed
+            existing_ids = set(collection.get(include=[])['ids'])
+        except Exception as e:
+            logger.warning(f"Could not retrieve existing IDs from ChromaDB: {e}. Assuming no existing IDs for now.")
+
+        new_chunks = [chunk for chunk in chunks if chunk.get('id') and chunk['id'] not in existing_ids]
+
         if not new_chunks:
-            print("No new chunks to add.")
+            logger.info("No new chunks to add.")
             return
-        batch_size = 128
+        
+        logger.info(f"Adding {len(new_chunks)} new chunks to the vector database...")
+
+        batch_size = 128 # Good batch size for embedding
         for i in range(0, len(new_chunks), batch_size):
             batch = new_chunks[i:i + batch_size]
             texts = [chunk['text'] for chunk in batch]
             ids = [chunk['id'] for chunk in batch]
-            metadatas = [self._flatten_metadata(chunk['metadata']) for chunk in batch]
+            metadatas = [self._flatten_metadata(chunk['metadata']) if chunk.get('metadata') else {} for chunk in batch]
+            
+            # Embed texts. This is the CPU-heavy part for BAAI/bge-large-en-v1.5
             embeddings = self.embedding_model.encode(texts, 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}")
+        logger.info(f"Finished adding {len(new_chunks)} chunks.")
 
     def search(self, query_text: str, top_k: int = None) -> List[Dict]:
         collection = self._get_collection()
+        
+        # Embed query text. This is also CPU-heavy for BAAI/bge-large-en-v1.5
         query_embedding = self.embedding_model.encode([query_text]).tolist()
+        
         top_k = top_k if top_k else self.top_k_default
         results = collection.query(
             query_embeddings=query_embedding,
             n_results=top_k,
             include=["documents", "metadatas", "distances"]
         )
+        
         search_results = []
-        for i, doc in enumerate(results['documents'][0]):
-            relevance_score = 1 - results['distances'][0][i]
-            search_results.append({
-                'text': doc,
-                'metadata': results['metadatas'][0][i],
-                'relevance_score': relevance_score
-            })
+        # Ensure results are not empty before accessing
+        if results and results['documents'] and results['documents'][0]:
+            for i, doc in enumerate(results['documents'][0]):
+                # ChromaDB distances are L2, which means smaller is better.
+                # Converting to a similarity score where 1 is perfect match, 0 is no match.
+                # A common conversion for L2 is 1 / (1 + distance) or max_dist - dist.
+                # Here, 1 - distance is used, assuming normalized embeddings leading to dist between 0 and 2.
+                relevance_score = 1 - results['distances'][0][i] 
+                search_results.append({
+                    'text': doc,
+                    'metadata': results['metadatas'][0][i],
+                    'relevance_score': relevance_score
+                })
         return search_results
 
 def ensure_db_populated(db_instance: PolicyVectorDB, chunks_file_path: str):
+    """
+    Ensures the ChromaDB is populated with data from the chunks file if it's currently empty.
+    """
     try:
+        # Check if the collection already has data
         if db_instance._get_collection().count() == 0:
+            logger.info("Vector database is empty. Attempting to populate from chunks file.")
             if not os.path.exists(chunks_file_path):
-                print(f"Chunks file not found at {chunks_file_path}")
+                logger.error(f"Chunks file not found at {chunks_file_path}. Cannot populate DB.")
                 return False
+            
             with open(chunks_file_path, 'r', encoding='utf-8') as f:
                 chunks_to_add = json.load(f)
+            
+            if not chunks_to_add:
+                logger.warning(f"Chunks file at {chunks_file_path} is empty. No data to add to DB.")
+                return False
+
             db_instance.add_chunks(chunks_to_add)
+            logger.info("Vector database population attempt complete.")
             return True
         else:
+            logger.info("Vector database already contains data. Skipping population.")
             return True
     except Exception as e:
-        print(f"DB Population Error: {e}")
-        return False
+        logger.error(f"DB Population Error: {e}", exc_info=True) # exc_info for full traceback
+        return False