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
+from typing import List, Dict, Optional
 from sentence_transformers import SentenceTransformer
 import chromadb
 from chromadb.config import Settings
@@ -11,23 +11,20 @@ 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 = 5, relevance_threshold: float = 0.5):
+    def __init__(self, persist_directory: str, top_k_default: int = 3, relevance_threshold: float = 0.35):
         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
 
@@ -43,12 +40,23 @@ 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()}
+        """
+        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
 
     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:
@@ -61,70 +69,122 @@ class PolicyVectorDB:
         if not chunks_with_ids:
             return
 
-        existing_ids = set(collection.get(ids=[str(c['id']) for c in chunks_with_ids])['ids'])
+        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'])
         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...")
-        
-        # Process in batches for efficiency
-        batch_size = 32 # Reduced batch size for potentially large embeddings
+        batch_size = 32
         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()
-            
+            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()
             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]:
+    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]:
         """
-        Searches the vector database for a given query text.
-        Returns a list of results filtered by a relevance threshold.
+        Searches the vector database. Optionally filter by section or by keywords in the financial_keywords metadata field.
+        Results are filtered by a relevance threshold.
         """
         collection = self._get_collection()
-        
-        # ✅ IMPROVEMENT: Add the recommended instruction prefix for BGE retrieval models.
+
+        # Add 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
-        
-        # 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"]
-        )
-        
+
+        # 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)
+
         search_results = []
         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:
+            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:
                     search_results.append({
                         'text': doc,
-                        'metadata': results['metadatas'][0][i],
+                        'metadata': metadata,
                         '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.
@@ -138,7 +198,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:
@@ -146,7 +206,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