Add custom weights to rag indexes

Author: asamal4

ols/src/rag_index/index_loader.py

@@ -9,9 +9,15 @@
 
 logger = logging.getLogger(__name__)
 
+
+SCORE_DILUTION_WEIGHT = 0.05
+SCORE_DILUTION_DEPTH = 2
+
+
 # delay import of llama_index dependencies
 BaseIndex = Any
 BaseRetriever = Any
+QueryFusionRetriever = Any
 
 
 # NOTE: Loading/importing something from llama_index bumps memory
@@ -43,6 +49,35 @@ def load_llama_index_deps() -> None:
     from llama_index.vector_stores.faiss import FaissVectorStore
 
 
+class QueryFusionRetrieverCustom(QueryFusionRetriever):
+    def __init__(self, **kwargs):
+        """Initialize custom query fusion class."""
+        super().__init__(**kwargs)
+
+        retrievers = kwargs.get("retrievers", None)
+        retriever_weights = kwargs.get("retriever_weights", None)
+        if not retriever_weights:
+            retriever_weights = [1.0] * len(retrievers)
+        self._custom_retriever_weights = retriever_weights
+
+    def _simple_fusion(self, results):
+        """Override internal method and apply weighted score."""
+        # Overriding one of the method is okay, we just need to add our custom logic.
+        all_nodes = {}
+        for i, nodes_with_scores in enumerate(results.values()):
+            for j, node_with_score in enumerate(nodes_with_scores):
+                node_index_id = f"{i}_{j}"
+                all_nodes[node_index_id] = node_with_score
+                # weighted_score = node_with_score.score * self._custom_retriever_weights[i]
+                # Uncomment above and delete below, if we decide weights to be set from config.
+                weighted_score = node_with_score.score * (
+                    1 - min(i, SCORE_DILUTION_DEPTH - 1) * SCORE_DILUTION_WEIGHT
+                )
+                all_nodes[node_index_id].score = weighted_score
+
+        return sorted(all_nodes.values(), key=lambda x: x.score or 0.0, reverse=True)
+
+
 class IndexLoader:
     """Load index from local file storage."""
 
@@ -146,12 +181,16 @@ def get_retriever(
             and self._retriever.similarity_top_k == similarity_top_k
         ):
             return self._retriever
-        retriever = QueryFusionRetriever(
-            [
+
+        # Note: we are using a custom retriever, based on our need
+        retriever = QueryFusionRetrieverCustom(
+            retrievers=[
                 index.as_retriever(similarity_top_k=similarity_top_k)
                 for index in self._indexes
             ],
             similarity_top_k=similarity_top_k,
+            retriever_weights=None,  # Setting as None, until this gets added to config
+            mode="simple",  # Don't modify this as we are adding our own logic
             num_queries=1,  # set this to 1 to disable query generation
             use_async=False,
             verbose=False,