updates to policy_vector_db

app/policy_vector_db.py

@@ -7,20 +7,32 @@ import chromadb
 from chromadb.config import Settings
 import logging
 
-logger = logging.getLogger("app")
+# --- Basic Logging Setup ---
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
 
 class PolicyVectorDB:
+    """
+    Manages the connection, population, and querying of a ChromaDB vector database 
+    for policy documents.
+    """
     def __init__(self, persist_directory: str, top_k_default: int = 5, relevance_threshold: float = 0.5):
         self.persist_directory = persist_directory
         self.client = chromadb.PersistentClient(path=persist_directory, settings=Settings(allow_reset=True))
         self.collection_name = "neepco_dop_policies"
-        # ✅ Use 'cuda' if a GPU is available for better performance
-        self.embedding_model = SentenceTransformer('BAAI/bge-large-en-v1.5', device='cpu')
-        self.collection = self._get_collection()
+        
+        # Using a powerful open-source embedding model.
+        # Change 'cpu' to 'cuda' if a GPU is available for significantly faster embedding.
+        self.embedding_model = SentenceTransformer('BAAI/bge-large-en-v1.5', device='cpu') 
+        
+        self.collection = None # Initialize collection as None for lazy loading
         self.top_k_default = top_k_default
         self.relevance_threshold = relevance_threshold
 
     def _get_collection(self):
+        """
+        Retrieves or creates the ChromaDB collection. Implements lazy loading.
+        """
         if self.collection is None:
             self.collection = self.client.get_or_create_collection(
                 name=self.collection_name,
@@ -29,9 +41,13 @@ class PolicyVectorDB:
         return self.collection
 
     def _flatten_metadata(self, metadata: Dict) -> Dict:
+        """Ensures all metadata values are strings, as required by some ChromaDB versions."""
         return {key: str(value) for key, value in metadata.items()}
 
     def add_chunks(self, chunks: List[Dict]):
+        """
+        Adds a list of chunks to the vector database, skipping any that already exist.
+        """
         collection = self._get_collection()
         if not chunks:
             logger.info("No chunks provided to add.")
@@ -47,23 +63,32 @@ class PolicyVectorDB:
         new_chunks = [chunk for chunk in chunks_with_ids if str(chunk.get('id')) not in existing_ids]
 
         if not new_chunks:
-            logger.info("All provided chunks already exist in the database.")
+            logger.info("All provided chunks already exist in the database. No new data to add.")
             return
         
         logger.info(f"Adding {len(new_chunks)} new chunks to the vector database...")
+        
+        # Process in batches for efficiency
         batch_size = 64
         for i in range(0, len(new_chunks), batch_size):
             batch = new_chunks[i:i + batch_size]
+            
             ids = [str(chunk['id']) for chunk in batch]
             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()
+            
             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]:
+        """
+        Searches the vector database for a given query text.
+        Returns a list of results filtered by a relevance threshold.
+        """
         collection = self._get_collection()
         query_embedding = self.embedding_model.encode([query_text]).tolist()
         k = top_k if top_k is not None else self.top_k_default
@@ -80,7 +105,6 @@ class PolicyVectorDB:
             for i, doc in enumerate(results['documents'][0]):
                 relevance_score = 1 - results['distances'][0][i]
                 
-                # ✅ RECOMMENDED CHANGE: Filter results internally based on the threshold
                 if relevance_score >= self.relevance_threshold:
                     search_results.append({
                         'text': doc,
@@ -88,21 +112,22 @@ class PolicyVectorDB:
                         'relevance_score': relevance_score
                     })
         
-        # Return the top k results *after* filtering
         return sorted(search_results, key=lambda x: x['relevance_score'], reverse=True)[:k]
 
-def ensure_db_populated(db_instance: PolicyVectorDB, chunks_file_path: str):
+def ensure_db_populated(db_instance: PolicyVectorDB, chunks_file_path: str) -> bool:
+    """
+    Checks if the DB is empty and populates it from a JSONL file if needed.
+    """
     try:
         if db_instance._get_collection().count() > 0:
             logger.info("Vector database already contains data. Skipping population.")
             return True
-        
+
         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}. Cannot populate DB.")
+            logger.error(f"Chunks file not found at '{chunks_file_path}'. Cannot populate DB.")
             return False
         
-        # ✅ CORRECTED CODE: Read the JSONL file line-by-line
         chunks_to_add = []
         with open(chunks_file_path, 'r', encoding='utf-8') as f:
             for line in f:
@@ -112,12 +137,12 @@ def ensure_db_populated(db_instance: PolicyVectorDB, chunks_file_path: str):
                     logger.warning(f"Skipping malformed line in chunks file: {line.strip()}")
         
         if not chunks_to_add:
-            logger.warning(f"Chunks file at {chunks_file_path} is empty or invalid. No data to add.")
+            logger.warning(f"Chunks file at '{chunks_file_path}' is empty or invalid. No data to add.")
             return False
 
         db_instance.add_chunks(chunks_to_add)
         logger.info("Vector database population attempt complete.")
         return True
     except Exception as e:
-        logger.error(f"DB Population Error: {e}", exc_info=True)
+        logger.error(f"An error occurred during DB population check: {e}", exc_info=True)
         return False