INTPYTHON-964 Add rerank capability to MongoDB VectorStore and Retrievers (#396)
INTPYTHON-964
## Summary
Adds native reranking (`$rerank`) support across
`MongoDBAtlasVectorSearch` and all
retrievers in `langchain-mongodb`. `$rerank` is a new MongoDB
aggregation stage
(public preview, Atlas MongoDB 8.3+) that re-ranks a set of candidate
documents
using the Voyage AI reranker API, improving result relevance beyond what
vector
similarity alone can provide.
`$rerank` is particularly valuable where embedding models are weak —
most notably
negation: a document saying "the filling is NOT between the bread"
scores high on
vector similarity for the query "filling between two slices of bread"
because the
embedding model focuses on token overlap, not meaning. The Voyage
reranker reads
the full sentence and demotes it correctly.
## Changes in this PR
### New pipeline helper — `pipelines.py`
`rerank_stage(query, path, num_docs_to_rerank, model=None)` builds the
`$rerank`
aggregation stage and a `$set` that writes both `score` (consumed
internally) and
`rerankScore` (surfaced in document metadata) from `{$meta: "score"}`
after
reranking. Marked for migration to `pymongo_search_utils` in a
follow-up.
### `MongoDBAtlasVectorSearch` — `vectorstores.py`
`similarity_search` and `similarity_search_with_score` accept three new
kwargs:
| kwarg | purpose |
|---|---|
| `rerank_path` | field or list of fields to rerank on; enables
`$rerank` when set |
| `rerank_model` | Voyage AI model name (e.g. `"rerank-2.5-lite"`); omit
for latest |
| `num_docs_to_rerank` | candidates passed to the reranker (max 1000);
defaults to `k` |
When `rerank_path` is set, the `$vectorSearch` limit is expanded to
`num_docs_to_rerank` so the reranker has enough candidates, and `$limit
k` is
appended after `$rerank` to trim back to the requested result count.
### Retrievers
All retrievers gain `rerank_path`, `rerank_model`, and
`num_docs_to_rerank`
(as class fields or kwargs), with `$rerank` injected at the appropriate
point in
each pipeline:
| Retriever | Injection point |
|---|---|
| `MongoDBAtlasFullTextSearchRetriever` | After `$search` + `$limit
n_to_rerank` |
| `MongoDBAtlasHybridSearchRetriever` | After `final_hybrid_stage` |
| `MongoDBAtlasParentDocumentRetriever` | After `$replaceRoot` (reranks
full parent docs, not child chunks) |
| `MongoDBAtlasSelfQueryRetriever` | No code change — rerank kwargs flow
through `search_kwargs` into `similarity_search` |
| `MongoDBGraphRAGRetriever` | After `$replaceRoot` in
`related_entities()`; reranks graph-traversal results against the
original query |
### Requirements
- MongoDB Atlas cluster running MongoDB 8.3+
- Native Reranking enabled in Atlas Project Settings
- Voyage AI API key configured in Atlas
> **Note on models:** as of the current public preview, only
`rerank-2.5-lite` is
> backed by real GPUs. Other model names (`rerank-2.5`, `rerank-2`,
`rerank-2-lite`)
> return a constant score of `0.5987` for all documents.
## Test Plan
### Unit tests (`tests/unit_tests/`)
`test_vectorstores.py` and `test_retrievers.py` gain pipeline-inspection
tests
using a `MockCollectionCapturePipeline` that records the pipeline passed
to
`aggregate()`. They verify:
- `rerank_stage()` produces the correct stage structure
- `$rerank` appears in the pipeline when `rerank_path` is set
- `$limit k` is appended when `num_docs_to_rerank > k`
- A clear `ValueError` is raised when `rerank_path` is set but no query
text is
available (non-autoembedding path called without `rerank_query`)
- No `$rerank` stage is present when `rerank_path` is not set (default
behaviour
unchanged)
### Integration tests (`tests/integration_tests/test_rerank.py`)
Run against a real Atlas cluster with Native Reranking enabled. Skipped
automatically when `MONGODB_URI` points to `localhost` / `127.0.0.1`.
The corpus is deliberately adversarial for vector search — every
document contains
the query tokens "filling" and "two slices of bread", but documents 2–5
do so in
a negating context. The key test is
`test_rerank_changes_ordering_vs_vector_search`, which asserts
concretely:
- Vector search top-1: French toast (`"two slices of bread but the
filling is NOT
between them"` — fooled by token overlap)
- Reranker top-1: club sandwich (only document that affirmatively
describes filling
between bread)
Other integration tests cover: `k` results returned, scores are positive
floats,
scores are descending, `rerankScore` appears in metadata and matches the
returned
score, `num_docs_to_rerank > k` correctly trims to `k`, and both
`MongoDBAtlasFullTextSearchRetriever` and
`MongoDBAtlasHybridSearchRetriever`
return correct top-1 results end-to-end.
## Checklist
### Checklist for Author
- [x] Did you update the changelog (if necessary)?
- [x] Is the intention of the code captured in relevant tests?
- [x] If there are new TODOs, has a related JIRA ticket been created?
- [x] Has a MongoDB Employee run [the patch build of this
PR](https://github.com/mongodb-labs/ai-ml-pipeline-testing?tab=readme-ov-file#running-a-patch-build-of-a-given-pr)?
### Checklist for Reviewer
- [x] Does the title of the PR reference a JIRA Ticket?
- [ ] Do you fully understand the implementation? (Would you be
comfortable explaining how this code works to someone else?)
---------
Signed-off-by: Casey Clements <casey.clements@mongodb.com> C
Casey Clements committed
c674fab15c932f3a1cb83910d9352a8b2f528a1d
Parent: 8e6c8dd
Committed by GitHub <noreply@github.com>
on 6/12/2026, 9:17:22 PM