Update app/policy_vector_db.py

app/policy_vector_db.py

@@ -40,39 +40,40 @@ class PolicyVectorDB:
         try:
             existing_ids = set(collection.get(include=[])['ids'])
         except Exception as e:
-            logger.warning(f"Could not retrieve existing IDs from ChromaDB: {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:
             logger.info("No new chunks to add.")
             return
         
+        logger.info(f"Adding {len(new_chunks)} new chunks to the vector database...")
         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.get('metadata', {})) for chunk in batch]
+            metadatas = [self._flatten_metadata(chunk['metadata']) if chunk.get('metadata') else {} 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()
         query_embedding = self.embedding_model.encode([query_text]).tolist()
-        top_k = top_k or self.top_k_default
-
+        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 = []
-        if results and results['documents'][0]:
+        if results and results['documents'] and results['documents'][0]:
             for i, doc in enumerate(results['documents'][0]):
-                relevance_score = 1 - results['distances'][0][i]
+                relevance_score = 1 - results['distances'][0][i] 
                 search_results.append({
                     'text': doc,
                     'metadata': results['metadatas'][0][i],
@@ -83,24 +84,24 @@ class PolicyVectorDB:
 def ensure_db_populated(db_instance: PolicyVectorDB, chunks_file_path: str):
     try:
         if db_instance._get_collection().count() == 0:
-            logger.info("Vector database is empty. Attempting to populate...")
+            logger.info("Vector database is empty. Attempting to populate from chunks file.")
             if not os.path.exists(chunks_file_path):
-                logger.error(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("Chunks file is empty.")
+                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("Database population complete.")
+            logger.info("Vector database population attempt complete.")
             return True
         else:
-            logger.info("Database already populated.")
+            logger.info("Vector database already contains data. Skipping population.")
             return True
     except Exception as e:
         logger.error(f"DB Population Error: {e}", exc_info=True)
-        return False
+        return False