RAGProcessor.py

from typing import List, Dict
from Configurator import AppConfig, RerankMethod


class RAGProcessor:
    """
    Post-processing component for RAG pipeline.
    Handles re-ranking, context assembly, and prompt construction.
    """
    
    def __init__(self, config: AppConfig):
        """
        Initialize the RAG processor.
        
        Args:
            config (AppConfig): Application configuration
        """
        self.config = config
    
    def rerank_chunks(self, chunks: List[Dict], query: str) -> List[Dict]:
        """
        Re-rank retrieved chunks based on the configured method.
        
        Args:
            chunks (List[Dict]): Retrieved chunks with scores
            query (str): Original user query
            
        Returns:
            List[Dict]: Re-ranked chunks
        """
        if not chunks:
            return chunks
        
        rerank_method = self.config.rag.rerank_method
        
        if rerank_method == RerankMethod.MMR:
            return self._rerank_mmr(chunks, query)
        elif rerank_method == RerankMethod.CROSS_ENCODER:
            return self._rerank_cross_encoder(chunks, query)
        elif rerank_method == RerankMethod.LLM_RERANK:
            return self._rerank_llm(chunks, query)
        else:
            # Default: return top chunks by score
            return sorted(chunks, key=lambda x: x.get("score", 0), reverse=True)[:self.config.rag.final_chunks]
    
    def _rerank_mmr(self, chunks: List[Dict], query: str) -> List[Dict]:
        """
        Re-rank using Maximal Marginal Relevance (MMR).
        Balances relevance and diversity.
        
        Args:
            chunks (List[Dict]): Retrieved chunks
            query (str): Original query
            
        Returns:
            List[Dict]: MMR re-ranked chunks
        """
        if len(chunks) <= self.config.rag.final_chunks:
            return chunks
        
        # Sort by score first
        chunks = sorted(chunks, key=lambda x: x.get("score", 0), reverse=True)
        
        selected_chunks = []
        remaining_chunks = chunks.copy()
        
        # Select first chunk (highest score)
        if remaining_chunks:
            selected_chunks.append(remaining_chunks.pop(0))
        
        # MMR selection for remaining chunks
        while len(selected_chunks) < self.config.rag.final_chunks and remaining_chunks:
            best_chunk = None
            best_mmr_score = -1
            
            for chunk in remaining_chunks:
                # Calculate MMR score: λ * relevance - (1-λ) * max_similarity
                relevance = chunk.get("score", 0)
                max_similarity = self._calculate_max_similarity(chunk, selected_chunks)
                mmr_score = self.config.rag.mmr_diversity_threshold * relevance - (1 - self.config.rag.mmr_diversity_threshold) * max_similarity
                
                if mmr_score > best_mmr_score:
                    best_mmr_score = mmr_score
                    best_chunk = chunk
            
            if best_chunk:
                selected_chunks.append(best_chunk)
                remaining_chunks.remove(best_chunk)
            else:
                break
        
        return selected_chunks
    
    def _rerank_cross_encoder(self, chunks: List[Dict], query: str) -> List[Dict]:
        """
        Re-rank using cross-encoder (placeholder implementation).
        TODO: Implement actual cross-encoder model.
        
        Args:
            chunks (List[Dict]): Retrieved chunks
            query (str): Original query
            
        Returns:
            List[Dict]: Cross-encoder re-ranked chunks
        """
        # For now, fall back to MMR
        return self._rerank_mmr(chunks, query)
    
    def _rerank_llm(self, chunks: List[Dict], query: str) -> List[Dict]:
        """
        Re-rank using LLM (placeholder implementation).
        TODO: Implement LLM-based re-ranking.
        
        Args:
            chunks (List[Dict]): Retrieved chunks
            query (str): Original query
            
        Returns:
            List[Dict]: LLM re-ranked chunks
        """
        # For now, fall back to MMR
        return self._rerank_mmr(chunks, query)
    
    def _calculate_max_similarity(self, chunk: Dict, selected_chunks: List[Dict]) -> float:
        """
        Calculate maximum similarity between a chunk and selected chunks.
        Simple text-based similarity for now.
        
        Args:
            chunk (Dict): Chunk to compare
            selected_chunks (List[Dict]): Already selected chunks
            
        Returns:
            float: Maximum similarity score
        """
        if not selected_chunks:
            return 0.0
        
        chunk_text = chunk.get("text", "").lower()
        max_similarity = 0.0
        
        for selected_chunk in selected_chunks:
            selected_text = selected_chunk.get("text", "").lower()
            similarity = self._calculate_text_similarity(chunk_text, selected_text)
            max_similarity = max(max_similarity, similarity)
        
        return max_similarity
    
    def _calculate_text_similarity(self, text1: str, text2: str) -> float:
        """
        Calculate simple text similarity based on word overlap.
        
        Args:
            text1 (str): First text
            text2 (str): Second text
            
        Returns:
            float: Similarity score between 0 and 1
        """
        words1 = set(text1.split())
        words2 = set(text2.split())
        
        if not words1 or not words2:
            return 0.0
        
        intersection = words1.intersection(words2)
        union = words1.union(words2)
        
        return len(intersection) / len(union) if union else 0.0
    
    def assemble_context(self, chunks: List[Dict]) -> str:
        """
        Assemble selected chunks into a single context block.
        
        Args:
            chunks (List[Dict]): Selected chunks after re-ranking
            
        Returns:
            str: Assembled context with source citations
        """
        if not chunks:
            return ""
        
        context_parts = []
        current_tokens = 0
        max_tokens = self.config.rag.context_max_tokens
        
        for i, chunk in enumerate(chunks):
            chunk_text = chunk.get("text", "")
            source = chunk.get("source", "unknown")
            index = chunk.get("index", 0)
            
            # Estimate tokens (rough approximation: words * 1.3)
            chunk_tokens = int(len(chunk_text.split()) * 1.3)
            
            # Check if adding this chunk would exceed token limit
            if current_tokens + chunk_tokens > max_tokens and context_parts:
                break
            
            # Add chunk with source citation
            citation = f"[Source: {source}, Chunk: {index}]"
            context_part = f"{chunk_text}\n{citation}"
            context_parts.append(context_part)
            current_tokens += chunk_tokens
        
        return "\n\n".join(context_parts)
    
    def construct_prompt(self, query: str, context: str, conversation_context: str = "") -> List[Dict]:
        """
        Construct the prompt for the LLM with system, context, and user messages.
        
        Args:
            query (str): User's original query
            context (str): Assembled context from retrieved chunks
            conversation_context (str): Previous conversation context (optional)
            
        Returns:
            List[Dict]: Messages for the LLM
        """
        # Build system message with conversation context if available
        system_content = self.config.rag.system_prompt
        if conversation_context:
            system_content = f"{system_content}\n\n{conversation_context}"
        
        system_message = {
            "role": "system",
            "content": system_content
        }
        
        context_message = {
            "role": "user",
            "content": f"Context:\n{context}\n\nQuery: {query}"
        }
        
        return [system_message, context_message]
    
    def extract_sources(self, chunks: List[Dict]) -> List[Dict]:
        """
        Extract source information from chunks for citation.
        
        Args:
            chunks (List[Dict]): Selected chunks
            
        Returns:
            List[Dict]: Source information for citations
        """
        sources = []
        
        for chunk in chunks:
            source_info = {
                "source": chunk.get("source", "unknown"),
                "chunk_index": chunk.get("index", 0),
                "snippet": chunk.get("text", "")[:200] + "..." if len(chunk.get("text", "")) > 200 else chunk.get("text", ""),
                "score": chunk.get("score", 0.0)
            }
            sources.append(source_info)
        
        return sources
    
    def estimate_tokens(self, text: str) -> int:
        """
        Estimate token count for text.
        
        Args:
            text (str): Text to estimate
            
        Returns:
            int: Estimated token count
        """
        return int(len(text.split()) * 1.3)