imp_query_data.py

import argparse
import logging
import os
import hashlib
import pickle
import time
from typing import List, Dict, Tuple
import numpy as np
from langchain_chroma import Chroma
from langchain.prompts import ChatPromptTemplate
from langchain_ollama import OllamaLLM
from langchain_community.retrievers import BM25Retriever
from langchain.retrievers import EnsembleRetriever
from langchain_core.documents import Document
from get_embedding_function import get_embedding_function
from sentence_transformers import CrossEncoder

CHROMA_PATH = "chroma1"
BM25_CACHE_PATH = "bm25_cache"
# PROMPT_TEMPLATE = """
# You are an expert assistant analyzing scientific documents about {domain}.

# Answer the question based ONLY on the following context. Each document is marked with its relevance level.

# {context}

# Instructions:
# 1. Answer based ONLY on the provided context
# 2. Cite specific documents when making claims using [Document X] notation
# 3. If information is incomplete or not found in the context, explicitly state this
# 4. Use technical terminology appropriately when discussing scientific concepts

# Question: {question}

# Answer:"""
PROMPT_TEMPLATE = """
You are an expert assistant analyzing technical documents about {domain}.

Use the following context documents to answer the question. Extract specific facts, numbers, and technical details from the provided content.

Context Documents:
{context}

Instructions:
1. Extract relevant information from the context documents above
2. Provide specific answers with concrete details (numbers, formulas, percentages)
3. Cite documents using [Document X] notation when referencing specific facts
4. If you find numerical values or technical specifications, include them in your answer
5. Synthesize information from multiple documents if needed
6. Only state "information not available" if you genuinely cannot find any relevant details

Question: {question}

Answer:"""

# Configure logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")

class HybridRAGSystem:
    def __init__(self, embedding_model="all-mpnet-base-v2", reranker_model=None):
        self.embedding_function = get_embedding_function(embedding_model)
        self.db = Chroma(persist_directory=CHROMA_PATH, embedding_function=self.embedding_function)
        
        # Initialize reranker if specified
        if reranker_model:
            self.reranker = CrossEncoder(reranker_model)
        else:
            # Use a lightweight reranker by default
            self.reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
        
        # Performance metrics storage
        self.metrics = {
            'retrieval_times': [],
            'reranking_times': [],
            'generation_times': [],
            'total_times': []
        }
    
    def reciprocal_rank_fusion(self, results_lists: List[List[Tuple[Document, float]]], k: int = 60) -> List[Document]:
        """
        Implement Reciprocal Rank Fusion for combining multiple ranked lists.
        This typically performs better than simple weighted combination.
        """
        fused_scores = {}
        doc_map = {}
        
        for result_list in results_lists:
            for rank, (doc, score) in enumerate(result_list):
                # Use chunk_id as unique identifier
                doc_id = doc.metadata.get('chunk_id', hashlib.md5(doc.page_content.encode()).hexdigest())
                
                if doc_id not in fused_scores:
                    fused_scores[doc_id] = 0
                    doc_map[doc_id] = doc
                
                # RRF formula: 1 / (k + rank)
                fused_scores[doc_id] += 1 / (k + rank + 1)
        
        # Sort by fused score
        sorted_docs = sorted(fused_scores.items(), key=lambda x: x[1], reverse=True)
        
        # Return documents in order
        return [doc_map[doc_id] for doc_id, _ in sorted_docs]
    
    def adaptive_query_weighting(self, query: str) -> Tuple[float, float]:
        """
        Dynamically adjust weights based on query characteristics.
        Returns (vector_weight, bm25_weight)
        """
        query_length = len(query.split())
        
        # Check for technical/scientific terms (customize based on your domain)
        technical_terms = ['protein', 'dna', 'rna', 'gene', 'cell', 'enzyme', 'molecular', 
                          'biological', 'chemical', 'reaction', 'synthesis', 'metabolism']
        has_technical = any(term in query.lower() for term in technical_terms)
        
        # Check if query is a question
        is_question = query.strip().endswith('?') or any(query.lower().startswith(q) for q in ['what', 'how', 'why', 'when', 'where'])
        
        # Adaptive weighting logic
        if query_length <= 3 and not has_technical:
            # Short, general queries benefit from semantic search
            return 0.7, 0.3
        elif has_technical and query_length > 5:
            # Long technical queries benefit from BM25
            return 0.3, 0.7
        elif is_question:
            # Questions often benefit from semantic understanding
            return 0.6, 0.4
        else:
            # Default balanced approach
            return 0.5, 0.5
    
    def create_structured_context(self, documents: List[Document], scores: List[float], query: str) -> str:
        """
        Create a structured context with relevance indicators and better formatting.
        """
        # Detect domain from documents
        domain_keywords = {
            'biology': ['cell', 'protein', 'dna', 'gene'],
            'chemistry': ['reaction', 'compound', 'molecule'],
            'medicine': ['disease', 'treatment', 'patient'],
            'general': []
        }
        
        domain = 'general'
        doc_text = ' '.join([doc.page_content[:200].lower() for doc in documents[:3]])
        for d, keywords in domain_keywords.items():
            if any(kw in doc_text for kw in keywords):
                domain = d
                break
        
        context_parts = []
        for i, (doc, score) in enumerate(zip(documents, scores)):
            # Determine relevance level based on score
            if i < 2:
                relevance = "HIGH"
            elif i < 4:
                relevance = "MEDIUM"
            else:
                relevance = "LOW"
            
            # Format each document
            context_parts.append(
                f"[Document {i+1} - Relevance: {relevance} - Score: {score:.3f}]\n"
                f"Source: {doc.metadata.get('source', 'Unknown')}, Page: {doc.metadata.get('page', 'N/A')}\n"
                f"Content: {doc.page_content}\n"
            )
        
        context = "\n---\n".join(context_parts)
        return context, domain
    
    def retrieve_and_rerank(self, query: str, initial_k: int = 20, final_k: int = 5) -> Tuple[List[Document], List[float]]:
        """
        Perform hybrid retrieval with reranking.
        """
        start_time = time.time()
        
        # Get adaptive weights
        vector_weight, bm25_weight = self.adaptive_query_weighting(query)
        logging.info(f"Using weights - Vector: {vector_weight:.2f}, BM25: {bm25_weight:.2f}")
        
        # Create retrievers
        vector_retriever = self.db.as_retriever(search_kwargs={"k": initial_k})
        bm25_retriever = get_or_create_bm25_retriever(self.db)
        
        if bm25_retriever:
            bm25_retriever.k = initial_k
            
            # Retrieve from both
            vector_results = vector_retriever.invoke(query)
            bm25_results = bm25_retriever.invoke(query)
            
            # Apply RRF instead of simple ensemble
            all_results = self.reciprocal_rank_fusion([
                [(doc, 1.0) for doc in vector_results],
                [(doc, 1.0) for doc in bm25_results]
            ])
            
            # Take top candidates for reranking
            candidates = all_results[:initial_k]
        else:
            # Fallback to vector search only
            candidates = vector_retriever.invoke(query)
        
        retrieval_time = time.time() - start_time
        self.metrics['retrieval_times'].append(retrieval_time)
        
        # Reranking stage
        rerank_start = time.time()
        if len(candidates) > 0:
            # Prepare pairs for reranking
            pairs = [[query, doc.page_content] for doc in candidates]
            
            # Get reranking scores
            rerank_scores = self.reranker.predict(pairs)
            
            # Sort by reranking scores
            ranked_results = sorted(zip(candidates, rerank_scores), key=lambda x: x[1], reverse=True)
            
            # Get top-k after reranking
            final_docs = [doc for doc, score in ranked_results[:final_k]]
            final_scores = [score for doc, score in ranked_results[:final_k]]
        else:
            final_docs = candidates[:final_k]
            final_scores = [1.0] * len(final_docs)
        
        rerank_time = time.time() - rerank_start
        self.metrics['reranking_times'].append(rerank_time)
        
        logging.info(f"Retrieval: {retrieval_time:.3f}s, Reranking: {rerank_time:.3f}s")
        
        return final_docs, final_scores
    
    def query(self, query_text: str, verbose: bool = True) -> Dict:
        """
        Main query method with performance tracking.
        """
        total_start = time.time()
        
        # Retrieve and rerank
        documents, scores = self.retrieve_and_rerank(query_text)
        
        # Create structured context
        context, domain = self.create_structured_context(documents, scores, query_text)
        
        # Generate response
        gen_start = time.time()
        prompt_template = ChatPromptTemplate.from_template(PROMPT_TEMPLATE)
        prompt = prompt_template.format(
            domain=domain,
            context=context,
            question=query_text
        )
        
        model = OllamaLLM(
            model="mistral",
            temperature=0.1,
            stream=False
        )
        
        response_text = model.invoke(prompt)
        gen_time = time.time() - gen_start
        self.metrics['generation_times'].append(gen_time)
        
        total_time = time.time() - total_start
        self.metrics['total_times'].append(total_time)
        
        # Prepare result
        result = {
            'answer': response_text,
            'documents': documents,
            'scores': scores,
            'metrics': {
                'retrieval_time': self.metrics['retrieval_times'][-1],
                'reranking_time': self.metrics['reranking_times'][-1],
                'generation_time': gen_time,
                'total_time': total_time
            }
        }
        
        if verbose:
            self._print_formatted_response(result, query_text)
        
        return result
    
    def _print_formatted_response(self, result: Dict, query: str):
        """Pretty print the response with metrics."""
        print(f"\n{'='*80}")
        print(f"Query: {query}")
        print(f"{'='*80}\n")
        print(f"Answer:\n{result['answer']}\n")
        print(f"{'-'*80}")
        print("Sources:")
        for i, (doc, score) in enumerate(zip(result['documents'], result['scores']), 1):
            print(f"{i}. {doc.metadata.get('source', 'Unknown')} "
                  f"(Page {doc.metadata.get('page', 'N/A')}) - "
                  f"Relevance: {score:.3f}")
        print(f"{'-'*80}")
        print(f"Performance Metrics:")
        print(f"  Retrieval: {result['metrics']['retrieval_time']:.3f}s")
        print(f"  Reranking: {result['metrics']['reranking_time']:.3f}s")
        print(f"  Generation: {result['metrics']['generation_time']:.3f}s")
        print(f"  Total: {result['metrics']['total_time']:.3f}s")
        print(f"{'='*80}\n")
    
    def get_performance_summary(self) -> Dict:
        """Get summary of performance metrics."""
        return {
            'avg_retrieval_time': np.mean(self.metrics['retrieval_times']) if self.metrics['retrieval_times'] else 0,
            'avg_reranking_time': np.mean(self.metrics['reranking_times']) if self.metrics['reranking_times'] else 0,
            'avg_generation_time': np.mean(self.metrics['generation_times']) if self.metrics['generation_times'] else 0,
            'avg_total_time': np.mean(self.metrics['total_times']) if self.metrics['total_times'] else 0,
            'total_queries': len(self.metrics['total_times'])
        }

# Keep the original function signatures for backward compatibility
def get_collection_hash(db):
    """Generate a hash based on the document collection state"""
    try:
        all_docs = db.get()
        hasher = hashlib.md5()
        for doc, meta in zip(all_docs["documents"], all_docs["metadatas"]):
            hasher.update(f"{meta.get('chunk_id', '')}-{doc[:100]}".encode())
        return hasher.hexdigest()
    except Exception as e:
        logging.warning(f"Error generating collection hash: {e}")
        return None

def get_or_create_bm25_retriever(db, force_rebuild=False):
    """Get cached BM25 retriever or create a new one if needed"""
    os.makedirs(BM25_CACHE_PATH, exist_ok=True)
    
    collection_hash = get_collection_hash(db)
    cache_file = os.path.join(BM25_CACHE_PATH, f"bm25_{collection_hash}.pkl")
    
    if os.path.exists(cache_file) and not force_rebuild:
        logging.info("Loading BM25 retriever from cache")
        try:
            with open(cache_file, 'rb') as f:
                return pickle.load(f)
        except Exception as e:
            logging.warning(f"Error loading cached BM25 retriever: {e}")
    
    logging.info("Building new BM25 retriever")
    try:
        all_docs = db.get()
        documents = [
            Document(page_content=doc, metadata=meta)
            for doc, meta in zip(all_docs["documents"], all_docs["metadatas"])
        ]
        
        bm25_retriever = BM25Retriever.from_documents(documents)
        bm25_retriever.k = 5
        
        with open(cache_file, 'wb') as f:
            pickle.dump(bm25_retriever, f)
        
        return bm25_retriever
    except Exception as e:
        logging.error(f"Error creating BM25 retriever: {e}")
        return None

def query_rag(query_text: str):
    """Legacy function for backward compatibility"""
    rag_system = HybridRAGSystem()
    result = rag_system.query(query_text)
    return result['answer']

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("query_text", type=str, help="The query text.")
    parser.add_argument("--no-rerank", action="store_true", help="Disable reranking")
    parser.add_argument("--benchmark", action="store_true", help="Run performance benchmark")
    args = parser.parse_args()
    
    if args.benchmark:
        # Run benchmark mode
        benchmark_queries = [
            "What is DNA?",
            "Explain the process of protein synthesis",
            "How does RNA differ from DNA?",
            "What are the main functions of mitochondria?",
            "Describe the structure of a cell membrane"
        ]
        
        rag_system = HybridRAGSystem()
        
        print("Running benchmark...")
        for query in benchmark_queries:
            rag_system.query(query, verbose=False)
        
        # Print performance summary
        summary = rag_system.get_performance_summary()
        print("\nPerformance Summary:")
        print(f"Total queries: {summary['total_queries']}")
        print(f"Average retrieval time: {summary['avg_retrieval_time']:.3f}s")
        print(f"Average reranking time: {summary['avg_reranking_time']:.3f}s")
        print(f"Average generation time: {summary['avg_generation_time']:.3f}s")
        print(f"Average total time: {summary['avg_total_time']:.3f}s")
    else:
        # Normal query mode
        rag_system = HybridRAGSystem()
        rag_system.query(args.query_text)

if __name__ == "__main__":
    main()