Update app/policy_vector_db.py

app/policy_vector_db.py

@@ -7,22 +7,17 @@ import chromadb
 from chromadb.config import Settings
 import logging
 
-logger = logging.getLogger("app") # Use the same logger as app.py
+logger = logging.getLogger("app")
 
 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 # Note: set to 0 in app.py
+        self.relevance_threshold = relevance_threshold
 
     def _get_collection(self):
         if self.collection is None:
@@ -33,8 +28,6 @@ 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]):
@@ -43,58 +36,43 @@ class PolicyVectorDB:
             logger.info("No chunks provided to add.")
             return
 
-        # 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.")
+            logger.warning(f"Could not retrieve existing IDs from ChromaDB: {e}")
 
         new_chunks = [chunk for chunk in chunks if chunk.get('id') and chunk['id'] not in existing_ids]
-
         if not new_chunks:
             logger.info("No new chunks to add.")
             return
         
-        logger.info(f"Adding {len(new_chunks)} new chunks to the vector database...")
-
-        batch_size = 128 # Good batch size for embedding
+        batch_size = 128
         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']) if chunk.get('metadata') else {} for chunk in batch]
-            
-            # Embed texts. This is the CPU-heavy part for BAAI/bge-large-en-v1.5
+            metadatas = [self._flatten_metadata(chunk.get('metadata', {})) for chunk in batch]
             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"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
+        top_k = top_k or self.top_k_default
+
         results = collection.query(
             query_embeddings=query_embedding,
             n_results=top_k,
             include=["documents", "metadatas", "distances"]
         )
-        
+
         search_results = []
-        # Ensure results are not empty before accessing
-        if results and results['documents'] and results['documents'][0]:
+        if results 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] 
+                relevance_score = 1 - results['distances'][0][i]
                 search_results.append({
                     'text': doc,
                     'metadata': results['metadatas'][0][i],
@@ -103,30 +81,26 @@ class PolicyVectorDB:
         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.")
+            logger.info("Vector database is empty. Attempting to populate...")
             if not os.path.exists(chunks_file_path):
-                logger.error(f"Chunks file not found at {chunks_file_path}. Cannot populate DB.")
+                logger.error(f"Chunks file not found at {chunks_file_path}")
                 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.")
+                logger.warning("Chunks file is empty.")
                 return False
 
             db_instance.add_chunks(chunks_to_add)
-            logger.info("Vector database population attempt complete.")
+            logger.info("Database population complete.")
             return True
         else:
-            logger.info("Vector database already contains data. Skipping population.")
+            logger.info("Database already populated.")
             return True
     except Exception as e:
-        logger.error(f"DB Population Error: {e}", exc_info=True) # exc_info for full traceback
-        return False
+        logger.error(f"DB Population Error: {e}", exc_info=True)
+        return False