ComfyUI/custom_nodes/red_ribbon/socialtoolkit/architecture/top10_document_retrieval.py

from pydantic import BaseModel
from typing import Dict, List, Any, Optional, Tuple
import logging
import numpy as np

logger = logging.getLogger(__name__)

class Top10DocumentRetrievalConfigs(BaseModel):
    """Configuration for Top-10 Document Retrieval workflow"""
    retrieval_count: int = 10  # Number of documents to retrieve
    similarity_threshold: float = 0.6  # Minimum similarity score
    ranking_method: str = "cosine_similarity"  # Options: cosine_similarity, dot_product, euclidean
    use_filter: bool = False  # Whether to filter results
    filter_criteria: Dict[str, Any] = {}
    use_reranking: bool = False  # Whether to use reranking

class Top10DocumentRetrieval:
    """
    Top-10 Document Retrieval system based on mermaid chart in README.md
    Performs vector search to find the most relevant documents
    """
    
    def __init__(self, resources: Dict[str, Any], configs: Top10DocumentRetrievalConfigs):
        """
        Initialize with injected dependencies and configuration
        
        Args:
            resources: Dictionary of resources including search services
            configs: Configuration for Top-10 Document Retrieval
        """
        self.resources = resources
        self.configs = configs
        
        # Extract needed services from resources
        self.encoder_service = resources.get("encoder_service")
        self.similarity_search_service = resources.get("similarity_search_service")
        self.document_storage = resources.get("document_storage_service")
        
        logger.info("Top10DocumentRetrieval initialized with services")

    def execute(self, 
                    input_data_point: str, 
                    documents: List[Any] = None, 
                    document_vectors: List[Any] = None
                   ) -> dict[str, Any]:
        """
        Execute the document retrieval flow based on the mermaid chart
        
        Args:
            input_data_point: The query or information request
            documents: Optional list of documents to search
            document_vectors: Optional list of document vectors to search
            
        Returns:
            Dictionary of documents containing potentially relevant documents, along with potentially relevant metadata.
        """
        logger.info(f"Starting top-10 document retrieval for: {input_data_point}")
        
        # Step 1: Encode the query
        encoded_query = self._encode_query(input_data_point)
        
        # Step 2: Get vector embeddings and document IDs from storage if not provided
        if documents is None or document_vectors is None:
            documents, document_vectors = self._get_documents_and_vectors()
        
        # Step 3: Perform similarity search
        similarity_scores, doc_ids = self._similarity_search(
            encoded_query, 
            document_vectors, 
            [doc.get("id") for doc in documents]
        )
        
        # Step 4: Rank and sort results
        ranked_results = self._rank_and_sort_results(similarity_scores, doc_ids)
        
        # Step 5: Filter to top-N results
        top_doc_ids = self._filter_to_top_n(ranked_results)
        
        # Step 6: Retrieve potentially relevant documents
        potentially_relevant_docs = self._retrieve_relevant_documents(documents, top_doc_ids)
        
        logger.info(f"Retrieved {len(potentially_relevant_docs)} potentially relevant documents")
        return {
            "relevant_documents": potentially_relevant_docs,
            "scores": {doc_id: score for doc_id, score in ranked_results},
            "top_doc_ids": top_doc_ids
        }

    def retrieve_top_documents(self, input_data_point: str, documents: List[Any], document_vectors: List[Any]) -> List[Any]:
        """
        Public method to retrieve top documents for an input query
        
        Args:
            input_data_point: The query to search for
            documents: Documents to search
            document_vectors: Vectors for the documents
            
        Returns:
            List of potentially relevant documents
        """
        result = self.control_flow(input_data_point, documents, document_vectors)
        return result["relevant_documents"]
    
    def _encode_query(self, input_data_point: str) -> Any:
        """
        Encode the input query into a vector representation
        
        Args:
            input_data_point: The query to encode
            
        Returns:
            Vector representation of the query
        """
        logger.debug(f"Encoding query: {input_data_point}")
        return self.encoder_service.encode(input_data_point)
    
    def _get_documents_and_vectors(self) -> Tuple[List[Any], List[Any]]:
        """
        Get all documents and their vectors from storage
        
        Returns:
            Tuple of (documents, document_vectors)
        """
        logger.debug("Getting documents and vectors from storage")
        return self.document_storage.get_documents_and_vectors()
    
    def _similarity_search(self, encoded_query: Any, document_vectors: List[Any], 
                          doc_ids: List[str]) -> Tuple[List[float], List[str]]:
        """
        Perform similarity search between the query and document vectors
        
        Args:
            encoded_query: Vector representation of the query
            document_vectors: List of document vector embeddings
            doc_ids: List of document IDs corresponding to the vectors
            
        Returns:
            Tuple of (similarity_scores, document_ids)
        """
        logger.debug("Performing similarity search")
        
        # In a real implementation, this would use an efficient vector search
        similarity_scores = []
        
        for vector in document_vectors:
            if self.configs.ranking_method == "cosine_similarity":
                score = self._cosine_similarity(encoded_query, vector.get("embedding"))
            elif self.configs.ranking_method == "dot_product":
                score = self._dot_product(encoded_query, vector.get("embedding"))
            elif self.configs.ranking_method == "euclidean":
                score = self._euclidean_distance(encoded_query, vector.get("embedding"))
                # Convert distance to similarity score (higher is more similar)
                score = 1.0 / (1.0 + score)
            else:
                score = 0.0
                
            similarity_scores.append(score)
        
        # If the similarity search service is available, use it instead
        if self.similarity_search_service:
            return self.similarity_search_service.search(
                encoded_query, document_vectors, doc_ids
            )
            
        return similarity_scores, doc_ids
    
    def _rank_and_sort_results(self, similarity_scores: List[float], 
                              doc_ids: List[str]) -> List[Tuple[str, float]]:
        """
        Rank and sort results by similarity score
        
        Args:
            similarity_scores: List of similarity scores
            doc_ids: List of document IDs
            
        Returns:
            List of (document_id, score) tuples sorted by score
        """
        logger.debug("Ranking and sorting results")
        
        # Create a list of (document_id, score) tuples
        result_tuples = list(zip(doc_ids, similarity_scores))
        
        # Sort by score in descending order
        sorted_results = sorted(result_tuples, key=lambda x: x[1], reverse=True)
        
        return sorted_results
    
    def _filter_to_top_n(self, ranked_results: List[Tuple[str, float]]) -> List[str]:
        """
        Filter to top N results
        
        Args:
            ranked_results: List of (document_id, score) tuples
            
        Returns:
            List of top N document IDs
        """
        logger.debug(f"Filtering to top {self.configs.retrieval_count} results")
        
        # Apply threshold filter if configured
        filtered_results = []
        
        if self.configs.use_filter:
            for doc_id, score in ranked_results:
                if score >= self.configs.similarity_threshold:
                    filtered_results.append(doc_id)
        else:
            filtered_results = [doc_id for doc_id, _ in ranked_results]
        
        # Return top N results
        return filtered_results[:self.configs.retrieval_count]
    
    def _retrieve_relevant_documents(self, documents: List[Any], top_doc_ids: List[str]) -> List[Any]:
        """
        Retrieve potentially relevant documents
        
        Args:
            documents: List of all documents
            top_doc_ids: List of top document IDs
            
        Returns:
            List of potentially relevant documents
        """
        logger.debug("Retrieving potentially relevant documents")
        
        # Create a map of document ID to document for faster lookup
        doc_map = {doc.get("id"): doc for doc in documents}
        
        # Retrieve documents by ID
        relevant_docs = []
        
        for doc_id in top_doc_ids:
            if doc_id in doc_map:
                relevant_docs.append(doc_map[doc_id])
        
        return relevant_docs
    
    def _cosine_similarity(self, vec1: List[float], vec2: List[float]) -> float:
        """Calculate cosine similarity between two vectors"""
        if not vec1 or not vec2:
            return 0.0
            
        try:
            # Convert to numpy arrays for efficient calculation
            vec1_np = np.array(vec1)
            vec2_np = np.array(vec2)
            
            # Calculate dot product
            dot = np.dot(vec1_np, vec2_np)
            
            # Calculate norms
            norm1 = np.linalg.norm(vec1_np)
            norm2 = np.linalg.norm(vec2_np)
            
            # Calculate cosine similarity
            similarity = dot / (norm1 * norm2)
            return float(similarity)
        except Exception as e:
            logger.error(f"Error calculating cosine similarity: {e}")
            return 0.0
            
    def _dot_product(self, vec1: List[float], vec2: List[float]) -> float:
        """Calculate dot product between two vectors"""
        if not vec1 or not vec2:
            return 0.0
            
        try:
            # Convert to numpy arrays for efficient calculation
            vec1_np = np.array(vec1)
            vec2_np = np.array(vec2)
            
            # Calculate dot product
            dot = np.dot(vec1_np, vec2_np)
            return float(dot)
        except Exception as e:
            logger.error(f"Error calculating dot product: {e}")
            return 0.0
            
    def _euclidean_distance(self, vec1: List[float], vec2: List[float]) -> float:
        """Calculate Euclidean distance between two vectors"""
        if not vec1 or not vec2:
            return float('inf')
            
        try:
            # Convert to numpy arrays for efficient calculation
            vec1_np = np.array(vec1)
            vec2_np = np.array(vec2)
            
            # Calculate Euclidean distance
            distance = np.linalg.norm(vec1_np - vec2_np)
            return float(distance)
        except Exception as e:
            logger.error(f"Error calculating Euclidean distance: {e}")
            return float('inf')