Update app/policy_vector_db.py

app/policy_vector_db.py

@@ -13,18 +13,18 @@ logger = logging.getLogger(__name__)
 
 class PolicyVectorDB:
     """
-    Manages the connection, population, and querying of a ChromaDB vector database
+    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"
         
         # 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')
+        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
@@ -44,22 +44,13 @@ class PolicyVectorDB:
 
     def _flatten_metadata(self, metadata: Dict) -> Dict:
         """Ensures all metadata values are strings, as required by some ChromaDB versions."""
-        def flatten_value(value):
-            if isinstance(value, dict):
-                return str(value)
-            elif isinstance(value, list):
-                return str(value)
-            else:
-                return str(value)
-        
-        return {key: flatten_value(value) for key, value in metadata.items()}
+        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.")
             return
@@ -67,7 +58,6 @@ class PolicyVectorDB:
         chunks_with_ids = [c for c in chunks if c.get('id')]
         if len(chunks) != len(chunks_with_ids):
             logger.warning(f"Skipped {len(chunks) - len(chunks_with_ids)} chunks that were missing an 'id'.")
-
         if not chunks_with_ids:
             return
 
@@ -77,23 +67,24 @@ class PolicyVectorDB:
         if not new_chunks:
             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 = 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]
+            
             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]
-
+            
             # 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}")
-
+            
         logger.info(f"Finished adding {len(new_chunks)} chunks.")
 
     def search(self, query_text: str, top_k: int = None) -> List[Dict]:
@@ -117,69 +108,23 @@ class PolicyVectorDB:
             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']:
-            for i, doc in enumerate(results['documents'][0]):  # Access first sublist
-                # The distance for normalized embeddings is often interpreted as 1 - cosine_similarity
-                relevance_score = 1 - results['distances'][0][i]  # ✅ Fixed: Access distances correctly
-                if relevance_score >= self.relevance_threshold:
-                    search_results.append({
-                        'text': doc,
-                        'metadata': results['metadatas'][0][i],  # ✅ Fixed: Access metadatas correctly
-                        '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 search_with_filters(self, query_text: str, top_k: int = None, 
-                           section_filter: str = None, chunk_type_filter: str = None) -> List[Dict]:
-        """Enhanced search with metadata filtering capability."""
-        collection = self._get_collection()
-        
-        instructed_query = f"Represent this sentence for searching relevant passages: {query_text}"
-        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
-        
-        # Build where clause for filtering
-        where_clause = {}
-        if section_filter:
-            where_clause["section"] = section_filter
-        if chunk_type_filter:
-            where_clause["chunk_type"] = chunk_type_filter
-        
-        try:
-            results = collection.query(
-                query_embeddings=query_embedding,
-                n_results=k * 2,
-                include=["documents", "metadatas", "distances"],
-                where=where_clause if where_clause else None
-            )
-        except Exception as e:
-            logger.warning(f"Filtered search failed, falling back to regular search: {e}")
-            # Fall back to regular search if filtering fails
-            results = collection.query(
-                query_embeddings=query_embedding,
-                n_results=k * 2,
-                include=["documents", "metadatas", "distances"]
-            )
         
         search_results = []
-        if results and results.get('documents') and results['documents']:
-            for i, doc in enumerate(results['documents'][0]):  # Access first sublist
-                relevance_score = 1 - results['distances'][0][i]  # ✅ Fixed: Access distances correctly
+        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:
                     search_results.append({
                         'text': doc,
-                        'metadata': results['metadatas'][0][i],  # ✅ Fixed: Access metadatas correctly
+                        'metadata': results['metadatas'][0][i],
                         '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:
     """
     Checks if the DB is empty and populates it from a JSONL file if needed.
@@ -190,11 +135,10 @@ def ensure_db_populated(db_instance: PolicyVectorDB, chunks_file_path: str) -> b
             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.")
             return False
-
+        
         chunks_to_add = []
         with open(chunks_file_path, 'r', encoding='utf-8') as f:
             for line in f:
@@ -202,7 +146,7 @@ def ensure_db_populated(db_instance: PolicyVectorDB, chunks_file_path: str) -> b
                     chunks_to_add.append(json.loads(line))
                 except json.JSONDecodeError:
                     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.")
             return False
@@ -210,7 +154,6 @@ def ensure_db_populated(db_instance: PolicyVectorDB, chunks_file_path: str) -> b
         db_instance.add_chunks(chunks_to_add)
         logger.info("Vector database population attempt complete.")
         return True
-
     except Exception as e:
         logger.error(f"An error occurred during DB population check: {e}", exc_info=True)
-        return False
+        return False