app/policy_vector_db.py

import os
import json
import torch
from typing import List, Dict
from sentence_transformers import SentenceTransformer
import chromadb
from chromadb.config import Settings
import logging

# --- 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"
        
        # 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') 
        logger.info("Embedding model loaded successfully.")
        
        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,
                metadata={"description": "NEEPCO Delegation of Powers Policy"}
            )
        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 expand_query(self, query_text: str) -> List[str]:
        """
        Generates query variations to improve retrieval.
        Uses simple heuristics - zero LLM cost.
        """
        queries = [query_text]
        
        # Expand with synonyms for policy-related terms
        synonyms = {
            "approval": ["approval", "consent", "authorization", "permission"],
            "limit": ["limit", "threshold", "ceiling", "maximum"],
            "authority": ["authority", "official", "person", "representative"],
            "delegate": ["delegate", "authorize", "empower", "assign"],
            "power": ["power", "authority", "delegation", "responsibility"],
            "financial": ["financial", "monetary", "funds", "budget"],
        }
        
        for term, variants in synonyms.items():
            if term in query_text.lower():
                for variant in variants:
                    if variant.lower() not in query_text.lower():
                        expanded = query_text.replace(term, variant)
                        if expanded not in queries:
                            queries.append(expanded)
                        if len(queries) >= 4:
                            break
            if len(queries) >= 4:
                break
        
        return queries[:4]  # Limit to 4 variations

    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

        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

        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...")
        
        # 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]:
        """
        Searches the vector database for a given query text with expansion.
        Returns a list of results filtered by a relevance threshold.
        """
        collection = self._get_collection()
        k = top_k if top_k is not None else self.top_k_default
        
        # Expand query for better recall
        queries = self.expand_query(query_text)
        all_results = {}
        
        for query in queries:
            # Add the recommended instruction prefix for BGE retrieval models.
            instructed_query = f"Represent this sentence for searching relevant passages: {query}"
            
            # Normalize embeddings for more accurate similarity search.
            query_embedding = self.embedding_model.encode([instructed_query], normalize_embeddings=True).tolist()
            
            # 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"]
            )
            
            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:
                        key = doc  # Use document text as key
                        # Keep highest relevance score for duplicate documents
                        if key not in all_results or relevance_score > all_results[key]['relevance_score']:
                            all_results[key] = {
                                'text': doc,
                                'metadata': results['metadatas'][0][i],
                                'relevance_score': relevance_score
                            }
        
        # Sort by relevance score and return the top_k results
        return sorted(all_results.values(), 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.
    """
    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.")
            return False
        
        chunks_to_add = []
        with open(chunks_file_path, 'r', encoding='utf-8') as f:
            for line in f:
                try:
                    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

        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