Update app/policy_vector_db.py

app/policy_vector_db.py

@@ -1,7 +1,7 @@
 import os
 import json
 import torch
-from typing import List, Dict, Optional
+from typing import List, Dict
 from sentence_transformers import SentenceTransformer
 import chromadb
 from chromadb.config import Settings
@@ -11,20 +11,23 @@ import logging
 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 = 3, relevance_threshold: float = 0.35):
+    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
+        
+        self.collection = None # Initialize collection as None for lazy loading
         self.top_k_default = top_k_default
         self.relevance_threshold = relevance_threshold
 
@@ -40,23 +43,12 @@ class PolicyVectorDB:
         return self.collection
 
     def _flatten_metadata(self, metadata: Dict) -> Dict:
-        """
-        Ensures all metadata values are strings, including flattening lists.
-        """
-        flattened = {}
-        for key, value in metadata.items():
-            if isinstance(value, list):
-                flattened[key] = ', '.join(map(str, value))
-            elif value is None:
-                flattened[key] = ''
-            else:
-                flattened[key] = str(value)
-        return flattened
+        """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.
-        Token count is added for each chunk if possible.
         """
         collection = self._get_collection()
         if not chunks:
@@ -69,122 +61,70 @@ class PolicyVectorDB:
         if not chunks_with_ids:
             return
 
-        ids_to_check = [str(c['id']) for c in chunks_with_ids]
-        existing_records = collection.get(ids=ids_to_check)
-        existing_ids = set(existing_records['ids'])
+        existing_ids = set(collection.get(ids=[str(c['id']) for c in chunks_with_ids])['ids'])
         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. No new data to add.")
             return
-
+        
         logger.info(f"Adding {len(new_chunks)} new chunks to the vector database...")
-        batch_size = 32
+        
+        # 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 = []
-            for chunk in batch:
-                metadata = self._flatten_metadata(chunk.get('metadata', {}))
-                # Add token count for debugging/monitoring
-                try:
-                    import tiktoken
-                    encoding = tiktoken.get_encoding("cl100k_base")
-                    token_count = len(encoding.encode(chunk['text']))
-                    metadata['token_count'] = str(token_count)
-                except Exception:
-                    pass
-                metadatas.append(metadata)
-
-            embeddings = self.embedding_model.encode(
-                texts, normalize_embeddings=True, show_progress_bar=False
-            ).tolist()
+            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: Optional[int] = None, filter_by_section: Optional[str] = None,
-               filter_by_keywords: Optional[List[str]] = None) -> List[Dict]:
+    def search(self, query_text: str, top_k: int = None) -> List[Dict]:
         """
-        Searches the vector database. Optionally filter by section or by keywords in the financial_keywords metadata field.
-        Results are filtered by a relevance threshold.
+        Searches the vector database for a given query text.
+        Returns a list of results filtered by a relevance threshold.
         """
         collection = self._get_collection()
-
-        # Add instruction prefix for BGE retrieval models.
+        
+        # ✅ IMPROVEMENT: Add the recommended instruction prefix for BGE retrieval models.
         instructed_query = f"Represent this sentence for searching relevant passages: {query_text}"
+        
+        # ✅ IMPROVEMENT: Normalize embeddings for more accurate similarity search.
         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
-
-        # Filtering logic
-        where_clause = {}
-        if filter_by_section:
-            where_clause["section"] = filter_by_section
-
-        query_params = {
-            "query_embeddings": query_embedding,
-            "n_results": k * 2,
-            "include": ["documents", "metadatas", "distances"]
-        }
-        if where_clause:
-            query_params["where"] = where_clause
-
-        results = collection.query(**query_params)
-
+        
+        # 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
+            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]
-                metadata = results['metadatas'][i]
-
-                # Keyword filtering if required
-                keep = True
-                if filter_by_keywords and metadata.get('financial_keywords'):
-                    stored_keywords = [kw.strip() for kw in metadata['financial_keywords'].split(',')]
-                    if not any(kw in stored_keywords for kw in filter_by_keywords):
-                        keep = False
-                if keep and relevance_score >= self.relevance_threshold:
+            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': metadata,
+                        '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 get_stats(self) -> Dict:
-        """
-        Get quick analytics for your DB: total chunks, unique sections, counts of keyworded chunks.
-        """
-        collection = self._get_collection()
-        total_count = collection.count()
-        if total_count == 0:
-            return {"total_chunks": 0}
-
-        sample = collection.get(limit=min(100, total_count), include=["metadatas"])
-        sections = set()
-        has_financial_keywords = 0
-        has_authority_keywords = 0
-
-        for metadata in sample.get('metadatas', []):
-            if metadata.get('section'):
-                sections.add(metadata['section'])
-            if metadata.get('financial_keywords'):
-                has_financial_keywords += 1
-            if metadata.get('authority_keywords'):
-                has_authority_keywords += 1
-
-        return {
-            "total_chunks": total_count,
-            "unique_sections": list(sections),
-            "chunks_with_financial_keywords": has_financial_keywords,
-            "chunks_with_authority_keywords": has_authority_keywords,
-        }
-
-
 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.
@@ -198,7 +138,7 @@ def ensure_db_populated(db_instance: PolicyVectorDB, chunks_file_path: str) -> b
         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:
@@ -206,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