Update app/policy_vector_db.py

app/policy_vector_db.py

@@ -13,9 +13,10 @@ 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))
@@ -24,10 +25,10 @@ class PolicyVectorDB:
         # 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
+        self.collection = None  # Initialize collection as None for lazy loading
         self.top_k_default = top_k_default
         self.relevance_threshold = relevance_threshold
 
@@ -44,13 +45,22 @@ class PolicyVectorDB:
 
     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 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()}
 
     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
@@ -58,6 +68,7 @@ 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
 
@@ -67,24 +78,23 @@ 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
+        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]:
@@ -105,26 +115,72 @@ class PolicyVectorDB:
         # Retrieve more results initially to allow for filtering
         results = collection.query(
             query_embeddings=query_embedding,
-            n_results=k * 2, # Retrieve more to filter by threshold
+            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'][0]:
-            for i, doc in enumerate(results['documents'][0]):
+            for i, doc in enumerate(results['documents']):
                 # 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],
                         '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'][0]:
+            for i, doc in enumerate(results['documents']):
+                relevance_score = 1 - results['distances'][i]
+                if relevance_score >= self.relevance_threshold:
+                    search_results.append({
+                        'text': doc,
+                        'metadata': results['metadatas'][0][i],
+                        'relevance_score': relevance_score
+                    })
+        
+        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.
@@ -135,10 +191,11 @@ 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:
@@ -146,7 +203,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
@@ -154,6 +211,7 @@ 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