feat(INTPYTHON-542): support weighted RRF (#168)

Author: s-xu-wafios

libs/langchain-mongodb/langchain_mongodb/pipelines.py

@@ -102,30 +102,34 @@ def combine_pipelines(
 
 
 def reciprocal_rank_stage(
-    score_field: str, penalty: float = 0, **kwargs: Any
+    score_field: str, penalty: float = 0, weight: float = 1, **kwargs: Any
 ) -> List[Dict[str, Any]]:
-    """Stage adds Reciprocal Rank Fusion weighting.
+    """
+    Stage adds Weighted Reciprocal Rank Fusion (WRRF) scoring.
 
-        First, it pushes documents retrieved from previous stage
-        into a temporary sub-document. It then unwinds to establish
-        the rank to each and applies the penalty.
+    First, it groups documents into an array, assigns rank by array index,
+    and then computes a weighted RRF score.
 
     Args:
-        score_field: A unique string to identify the search being ranked
-        penalty: A non-negative float.
-        extra_fields: Any fields other than text_field that one wishes to keep.
+        score_field: A unique string to identify the search being ranked.
+        penalty: A non-negative float (e.g., 60 for RRF-60). Controls the denominator.
+        weight: A float multiplier for this source's importance.
+        **kwargs: Ignored; allows future extensions or passthrough args.
 
     Returns:
-        RRF score
+        Aggregation pipeline stage for weighted RRF scoring.
     """
 
-    rrf_pipeline = [
+    return [
         {"$group": {"_id": None, "docs": {"$push": "$$ROOT"}}},
         {"$unwind": {"path": "$docs", "includeArrayIndex": "rank"}},
         {
             "$addFields": {
                 f"docs.{score_field}": {
-                    "$divide": [1.0, {"$add": ["$rank", penalty, 1]}]
+                    "$multiply": [
+                        weight,
+                        {"$divide": [1.0, {"$add": ["$rank", penalty, 1]}]},
+                    ]
                 },
                 "docs.rank": "$rank",
                 "_id": "$docs._id",
@@ -134,8 +138,6 @@ def reciprocal_rank_stage(
         {"$replaceRoot": {"newRoot": "$docs"}},
     ]
 
-    return rrf_pipeline  # type: ignore
-
 
 def final_hybrid_stage(
     scores_fields: List[str], limit: int, **kwargs: Any

libs/langchain-mongodb/langchain_mongodb/retrievers/hybrid_search.py

@@ -43,6 +43,10 @@ class MongoDBAtlasHybridSearchRetriever(BaseRetriever):
     """Penalty applied to vector search results in RRF: scores=1/(rank + penalty)"""
     fulltext_penalty: float = 60.0
     """Penalty applied to full-text search results in RRF: scores=1/(rank + penalty)"""
+    vector_weight: float = 1.0
+    """Weight applied to vector search results in RRF: score = weight * (1 / (rank + penalty + 1))"""
+    fulltext_weight: float = 1.0
+    """Weight applied to full-text search results in RRF: score = weight * (1 / (rank + penalty + 1))"""
     show_embeddings: float = False
     """If true, returned Document metadata will include vectors."""
     top_k: Annotated[
@@ -95,7 +99,11 @@ def _get_relevant_documents(
                 oversampling_factor=self.oversampling_factor,
             )
         ]
-        vector_pipeline += reciprocal_rank_stage("vector_score", self.vector_penalty)
+        vector_pipeline += reciprocal_rank_stage(
+            score_field="vector_score",
+            penalty=self.vector_penalty,
+            weight=self.vector_weight,
+        )
 
         combine_pipelines(pipeline, vector_pipeline, self.collection.name)
 
@@ -109,7 +117,11 @@ def _get_relevant_documents(
         )
 
         text_pipeline.extend(
-            reciprocal_rank_stage("fulltext_score", self.fulltext_penalty)
+            reciprocal_rank_stage(
+                score_field="fulltext_score",
+                penalty=self.fulltext_penalty,
+                weight=self.fulltext_weight,
+            )
         )
 
         combine_pipelines(pipeline, text_pipeline, self.collection.name)