Weighted RRF

[Shanorino]

In MongoDBAtlasHybridSearchRetriever, currently it's not possible to add weights separately to 2 retrievers (in other words, it's always 50-50).

A workaround is to change vector_penalty and fulltext_penalty. However, tuning the RRF parameter k changes the shape of the reciprocal rank function, while multiplying by a weight is a simple linear scaling of the entire curve. These two approaches are not mathematically equivalent and lead to different effects on how each rank contributes to the final score.

One thing to mention is that EnsembleRetriever in Langchain also implemented the weights (instead of changing the penalty constant k).

Thus, I'd propose implementing weights to the class MongoDBAtlasHybridSearchRetriever by adding 2 more stages to the Mongo pipeline, so that it can be instantiated this way:

retriever = MongoDBAtlasHybridSearchRetriever(
vectorstore = vector_store,
search_index_name = "search_index",
top_k = 5,
fulltext_weight=0.3,
vector_weight=0.7,
fulltext_penalty = 50,
vector_penalty = 50
)

[blink1073]

Hi @Shanorino, I opened https://jira.mongodb.org/browse/INTPYTHON-542 to track this issue.

[caseyclements]

Hi @Shanorino. Would you mind going into further detail on what changes you're suggesting? Maybe you can write down mathematically, or include a link?

[s-xu-wafios]

Hi @caseyclements , so if the original RRF can be represented this way:

$$\text{RRF}(d) = \sum_{i=1}^{n} \frac{1}{k + r_i(d)}$$

The weighted RRF would be like this:

$$\text{WeightedRRF}(d) = \sum_{i=1}^{n} w_i \cdot \frac{1}{k + r_i(d)}$$

That said, I added a "weighted score stage", where each retriever has a $multiply with weight

    def _weighted_score_stage(self, score_field: str, weight: float, penalty: float) -> List[Dict[str, Any]]:
        """
        Apply weighted scoring with custom penalty for either vector or full-text scores.
        """
        return [
            {"$group": {"_id": None, "docs": {"$push": "$$ROOT"}}},
            {"$unwind": {"path": "$docs", "includeArrayIndex": "rank"}},
            {
                "$addFields": {
                    f"docs.{score_field}": {
                        "$multiply": [
                            weight,
                            {"$divide": [1.0, {"$add": ["$rank", penalty, 1]}]},
                        ]
                    },
                    "_id": "$docs._id",
                }
            },
            {"$replaceRoot": {"newRoot": "$docs"}},
        ]

[caseyclements]

@s-xu-wafios This is great. Sorry for the delay responding. In my work building tests, I'd come to the conclusion that the RRF is missing a parameterization like the one you suggest, and think the changes you suggest would do a lot to handle many cases.. Would you like to open a PR?

[s-xu-wafios]

@caseyclements sure, actually we have extended the class and it's been working fine for us. I'd merge the added "_weighted_score_stage" to pipeline.py and adapt the retriever with parameterization and open a PR

[caseyclements]

@s-xu-wafios Perfect! If you open up a pull-request with your changes, we'll make a release after it's merged.