fix(memory): give the reranker the full candidate pool, not the truncated top-k by OfficialAbhinavSingh · Pull Request #6449 · mem0ai/mem0

OfficialAbhinavSingh · 2026-07-20T12:55:18Z

Linked Issue

Description

When search(rerank=True) is used, the reranker only received the already-truncated top-limit results, so reranking could reorder the final results but never surface a relevant memory the first-stage scorer ranked below limit — the main reason to run a reranker.

_search_vector_store over-fetches a candidate pool (internal_limit = max(limit * 4, 60)) for hybrid scoring, but score_and_rank(..., top_k=limit) truncates to exactly limit before returning. So the list passed to reranker.rerank(...) was already the final top-limit set, and the over-fetched pool was discarded one stage too early.

Concrete example: search("refund policy", filters={"user_id": "u1"}, top_k=5, rerank=True) where the bi-encoder ranks the relevant memory 7th (exactly what cross-encoders exist to fix). score_and_rank truncates to 5, dropping it; the reranker reorders 5 less-relevant candidates and never sees the correct one. Recall is identical with and without rerank=True.

This PR retrieves the broader candidate pool when reranking is enabled, then lets the reranker narrow it back to limit. Non-rerank search is unchanged (retrieval_limit == limit). If the rerank call fails, the over-fetched pool is trimmed back to limit so the limit contract still holds. Applied to both the sync Memory.search() and async AsyncMemory.search().

retrieval_limit = max(limit * 4, 60) if (rerank and self.reranker) else limit
original_memories = self._search_vector_store(
    query, effective_filters, retrieval_limit, threshold, ...
)
if rerank and self.reranker and original_memories:
    try:
        original_memories = self.reranker.rerank(query, original_memories, limit)
    except Exception as e:
        logger.warning(f"Reranking failed, using original results: {e}")
        original_memories = original_memories[:limit]

Type of Change

Bug fix (non-breaking change that fixes an issue)
New feature (non-breaking change that adds functionality)
Breaking change (fix or feature that would cause existing functionality to change)
Refactor (no functional changes)
Documentation update

Breaking Changes

N/A — non-rerank search is byte-for-byte unchanged. With reranking enabled, results are the same size (limit) but ranking quality improves, since the reranker now scores the full candidate pool.

Test Coverage

I added/updated unit tests
I added/updated integration tests
I tested manually (describe below)
No tests needed (explain why)

Added three tests in tests/memory/test_main.py:

test_search_rerank_uses_full_candidate_pool — asserts the reranker receives more than limit candidates (the pool) and is asked to narrow to limit; final results respect limit. Fails on main (reranker gets exactly limit) and passes with this change.
test_async_search_rerank_uses_full_candidate_pool — async counterpart.
test_search_rerank_failure_still_respects_limit — on rerank failure, the over-fetched pool is trimmed back to limit.

Full tests/memory/test_main.py suite (51 tests) and tests/rerankers/ pass; ruff check is clean on the changed source.

Checklist

My code follows the project's style guidelines
I have performed a self-review of my code
I have added tests that prove my fix/feature works
New and existing tests pass locally
I have updated documentation if needed

kartik-mem0

Changes requested

Fixes the real retrieve-then-rerank ordering bug from #6448, but the fix itself introduces a hot-path regression: the widened pool gets widened again.

To unblock:

mem0/memory/main.py:1455 (and the async twin) - stop double-widening before it reaches _search_vector_store (detail in the inline comment)

How I verified

Problem is real: reproduces on main at mem0/memory/main.py:1598-1655 (sync), :3241-3293 (async) - live sandbox, real Qdrant/pgvector, reranker saw only 5 candidates (== top_k=5).
Tests / red-green: pytest tests/memory/test_main.py -k rerank -> 3 passed. Reverting the hunk fails 2/3 tests; a surgical revert of just the failure-trim line fails the third (assert 60 == 5). All 3 are load-bearing.
CI: Python SDK (3.10/3.11/3.12), changelog, license/cla, CI Gate all green. Vercel red, unrelated docs-deploy gate.
Security: cleared, no new attack surface, pure arithmetic on an internal query-depth integer.
Bloat (ponytail): lean already, comment is proportionate to a subtle ordering bug.
Not verified: TS SDK (mem0-ts/src/oss/src/memory/index.ts) has the identical bug, neither PR here touches it - flagged as a follow-up, not a blocker.

Good catch on the actual bug and clean tests. Push the fix for the double-widening and I will re-review.
Open PR #6455 fixes the same root cause with an identical mem0/memory/main.py hunk but bundles an unrelated, untested mem0/embeddings/gemini.py change; this PR is the cleaner of the two once the item above is resolved.

kartik-mem0 · 2026-07-23T11:12:04Z

+        # The reranker then narrows the pool back down to `limit`. Without this,
+        # score_and_rank truncates to `limit` before the reranker runs, so
+        # reranking could only reorder the final results, never improve recall.
+        retrieval_limit = max(limit * 4, 60) if (rerank and self.reranker) else limit


This widened retrieval_limit (already max(limit*4, 60)) gets passed into _search_vector_store as its own limit parameter, and _search_vector_store independently recomputes internal_limit = max(limit*4, 60) off that already-widened number. With top_k=5 and reranking on, the real vector-store query depth compounds to 240 instead of the intended 60 - 4x deeper than needed, on every dense and keyword-search call.

Move the widen decision fully inside _search_vector_store (derive internal_limit once from the true limit, and truncate score_and_rank's top_k to it when the caller signals reranking) so search() does not pre-inflate a number _search_vector_store inflates again.

Verified against a58558b1: the compounding double-widening finding is fixed (mem0/memory/main.py:1622,1627,1629,1634 sync; :3276,3281,3283,3288 async). internal_limit is now derived once from the original limit, and return_k (not limit) controls only the final score_and_rank truncation, so retrieval depth for both dense and keyword search no longer compounds. Red-green: reverting to the pre-follow-up commit reproduces passed_limit == 60 instead of 5; 4/4 sandbox lane by mode runs pass on this head.

Review follow-up on mem0ai#6449. The prior fix pre-widened the caller's limit to max(limit*4, 60) and passed that into _search_vector_store, which then independently recomputes internal_limit = max(limit*4, 60) off the already widened number — so with top_k=5 and reranking on, the vector/keyword fetch ballooned to max(60*4, 60) = 240 rows instead of the intended 60. Pass the original limit and a rerank_pool flag instead. _search_vector_store keeps its single over-fetch (internal_limit off the original limit) and only changes its return count to that full pool when rerank_pool is set, so the reranker still sees the whole candidate pool while the DB fetch stays at one max(limit*4, 60). Regression guard added: the rerank tests now assert _search_vector_store is called with the original limit, not a pre-widened one.

OfficialAbhinavSingh · 2026-07-23T17:21:57Z

@kartik-mem0 Good catch on the double-widening — fixed in b8280f43.

Root cause was exactly as you described: the caller pre-widened limit to max(limit*4, 60) and passed that into _search_vector_store, which then recomputes internal_limit = max(limit*4, 60) off the already-widened number → max(60*4, 60) = 240 rows fetched for top_k=5 instead of the intended 60.

Fix: the caller now passes the original limit plus a rerank_pool flag (both sync and async). _search_vector_store keeps its single over-fetch (internal_limit computed from the original limit) and only changes its return count to that full pool when rerank_pool is set — so the reranker still sees the whole candidate pool while the DB fetch stays at one max(limit*4, 60). No behavior change on the non-rerank path.

Added a regression guard: the rerank tests now assert _search_vector_store is called with the original limit (not a pre-widened one) — I verified it's load-bearing (reintroducing the pre-widening turns it red). pytest -k rerank → 3 passed, full tests/memory/ → 355 passed, rerankers suite → 104 passed, ruff clean.

On the two follow-ups you flagged: happy to open the TS twin (mem0-ts has the identical rerank-truncation) as a separate PR if useful, and I agree #6455 bundling the untested gemini.py change makes this the cleaner vehicle for the core fix. Ready for re-review.

Review follow-up on mem0ai#6449. The prior fix pre-widened the caller's limit to max(limit*4, 60) and passed that into _search_vector_store, which then independently recomputes internal_limit = max(limit*4, 60) off the already widened number — so with top_k=5 and reranking on, the vector/keyword fetch ballooned to max(60*4, 60) = 240 rows instead of the intended 60. Pass the original limit and a rerank_pool flag instead. _search_vector_store keeps its single over-fetch (internal_limit off the original limit) and only changes its return count to that full pool when rerank_pool is set, so the reranker still sees the whole candidate pool while the DB fetch stays at one max(limit*4, 60). Regression guard added: the rerank tests now assert _search_vector_store is called with the original limit, not a pre-widened one.

kartik-mem0

LGTM, ready for squash-merge

Verified against a58558b1 (new head): the confirmed compounding 4x over-fetch is fixed. _search_vector_store now derives internal_limit once from the true limit, and a rerank_pool flag decides only the final score_and_rank truncation (mem0/memory/main.py:1622,1627).

Before you merge: nothing, this is self-contained.
Follow-up: TS side (mem0-ts/src/oss/src/memory/index.ts:1406,1535,1587) still has the original truncate-before-rerank bug from #6448, not the double-widening variant (that shape needs Python's two-layer split). Worth #6537 picking up the core fix.
Not verified: live reranker API calls; SpyReranker plus real Qdrant/pgvector only.

Risk: low, internal fetch/truncate split only, no public API change · Your time: ~5 min

How I verified

Problem is real (cycle 1): reproduced on main at mem0/memory/main.py:1655 (score_and_rank truncates to limit before the reranker ever sees the pool).
This cycle's delta reproduced pre-fix: reverted to the pre-follow-up commit (58727af5), kept PR-head tests: assert passed_limit == FINAL_LIMIT failed with 60 == 5, the exact compounded double-widen value.
Tests: hatch -e dev_py_3_11 run pytest tests/memory/test_main.py -k rerank -v -> 3 passed. Full module: pytest tests/memory/test_main.py -q -> 51 passed. ruff check clean.
Red-green: reverted, 2/3 rerank tests failed; restored, 3/3 pass. ✓
Sandbox (live): issue6448_rerank_truncated_pool, 4/4 lane×mode PASS (oss-qdrant, oss-pgvector × sync, async) against this head, real Qdrant/pgvector, no mocks.
CI: all required checks green; Vercel red (auth-gated preview deploy, non-blocking).
Security: cleared, no new input surface, rerank_pool is an internal bool derived from existing config.
Bloat (ponytail): lean already, one bool param plus one local, no new abstraction.
Linked issue: #6448

…ated topK TypeScript twin of mem0ai#6449. Memory.search() over-fetches a candidate pool (internalLimit = max(topK*4, 60)) but scoreAndRank truncates the ranked output to topK before the reranker runs, so the reranker only ever sees the already-truncated top-topK set. It can reorder those but can never surface a relevant memory the first-stage scorer ranked at topK+1..N, defeating the two-stage retrieve-then-rerank design (recall ends up identical with and without reranking). When a reranker will run, return the full over-fetched pool from scoreAndRank instead of truncating to topK, let the reranker narrow it back to topK, and trim to topK if the rerank call fails so a failure doesn't leak the whole pool. The DB fetch is not re-widened — internalLimit is still computed once from topK. Closes mem0ai#6536

…ated top-k search(rerank=True) passed the reranker the already-truncated top-`limit` results, because score_and_rank truncates to `limit` before returning. The reranker could then only reorder the final results and never surface a relevant memory the first-stage scorer ranked below `limit` -- defeating the two-stage retrieve-then-rerank design. Retrieve a larger candidate pool (max(limit * 4, 60)) when reranking is enabled and let the reranker narrow it back to `limit`; trim the pool back to `limit` if the rerank call fails. Applied to both sync and async search(); non-rerank search is unchanged. Closes mem0ai#6448

Review follow-up on mem0ai#6449. The prior fix pre-widened the caller's limit to max(limit*4, 60) and passed that into _search_vector_store, which then independently recomputes internal_limit = max(limit*4, 60) off the already widened number — so with top_k=5 and reranking on, the vector/keyword fetch ballooned to max(60*4, 60) = 240 rows instead of the intended 60. Pass the original limit and a rerank_pool flag instead. _search_vector_store keeps its single over-fetch (internal_limit off the original limit) and only changes its return count to that full pool when rerank_pool is set, so the reranker still sees the whole candidate pool while the DB fetch stays at one max(limit*4, 60). Regression guard added: the rerank tests now assert _search_vector_store is called with the original limit, not a pre-widened one.

…ated topK TypeScript twin of mem0ai#6449. Memory.search() over-fetches a candidate pool (internalLimit = max(topK*4, 60)) but scoreAndRank truncates the ranked output to topK before the reranker runs, so the reranker only ever sees the already-truncated top-topK set. It can reorder those but can never surface a relevant memory the first-stage scorer ranked at topK+1..N, defeating the two-stage retrieve-then-rerank design (recall ends up identical with and without reranking). When a reranker will run, return the full over-fetched pool from scoreAndRank instead of truncating to topK, let the reranker narrow it back to topK, and trim to topK if the rerank call fails so a failure doesn't leak the whole pool. The DB fetch is not re-widened — internalLimit is still computed once from topK. Closes mem0ai#6536

OfficialAbhinavSingh · 2026-08-09T11:29:34Z

@kartik-mem0 checking in on this one. It has been sitting since your "ready for squash-merge" on 24 Jul.

I rebased it onto main on 4 Aug (2.0.13 to 2.0.15) to keep changelog_check from going stale, and CI Gate is green on the new head 6766827d. The diff is unchanged from the a58558b1 you verified.

Anything outstanding on my side? The TypeScript twin #6537, which you flagged as the follow-up in your review, is open and still awaiting a first review.

kartik-mem0

Re-reviewed at 6766827d. Content-identical rebase of the previously reviewed head. Fresh tests and lint pass against the new SHA. LGTM still holds.

OfficialAbhinavSingh · 2026-09-21T11:05:22Z

Thanks for re-verifying both of these today.

One thing that may not be visible from your side: every review on this PR and on #6516 is COMMENTED (4 each, none APPROVED), and .github/AGENTS.md describes Main Branch Rule as requiring one approving review with no bypass actors. So both still show "Review required" under the LGTM and sit at mergeable_state=blocked.

Is an explicit approval (or a direct merge) what's needed here, or is there something still outstanding on my side?

kartik-mem0 mentioned this pull request Jul 23, 2026

search(rerank=True) reranks only the truncated top-k, so reranking can't improve recall #6448

Open

kartik-mem0 reviewed Jul 23, 2026

View reviewed changes

kartik-mem0 mentioned this pull request Jul 23, 2026

fix(memory): fetch full candidate pool for reranker and fix tests #6455

Open

9 tasks

github-actions Bot added the sdk-python Python SDK specific label Jul 23, 2026

This was referenced Jul 23, 2026

bug(ts-oss): reranker only sees the truncated topK, not the full candidate pool #6536

Open

fix(ts-oss): give the reranker the full candidate pool, not the truncated topK #6537

Open

OfficialAbhinavSingh force-pushed the fix/rerank-candidate-pool branch from b8280f4 to a58558b Compare July 23, 2026 17:35

kartik-mem0 reviewed Jul 24, 2026

View reviewed changes

abhay-codes07 mentioned this pull request Jul 24, 2026

fix(ts-oss): feed the reranker the full candidate pool, not the truncated topK #6576

Closed

6 tasks

OfficialAbhinavSingh added 2 commits August 4, 2026 20:52

OfficialAbhinavSingh force-pushed the fix/rerank-candidate-pool branch from a58558b to 6766827 Compare August 4, 2026 15:25

kartik-mem0 reviewed Sep 21, 2026

View reviewed changes

Navigation Menu

fix(memory): give the reranker the full candidate pool, not the truncated top-k - #6449

fix(memory): give the reranker the full candidate pool, not the truncated top-k#6449
OfficialAbhinavSingh wants to merge 2 commits into
mem0ai:mainfrom
OfficialAbhinavSingh:fix/rerank-candidate-pool

OfficialAbhinavSingh commented Jul 20, 2026

Uh oh!

kartik-mem0 left a comment

Uh oh!

kartik-mem0 Jul 23, 2026

Uh oh!

kartik-mem0 Jul 24, 2026

Uh oh!

OfficialAbhinavSingh commented Jul 23, 2026

Uh oh!

kartik-mem0 left a comment

Uh oh!

OfficialAbhinavSingh commented Aug 9, 2026

Uh oh!

kartik-mem0 left a comment

Uh oh!

OfficialAbhinavSingh commented Sep 21, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

OfficialAbhinavSingh commented Jul 20, 2026

Linked Issue

Description

Type of Change

Breaking Changes

Test Coverage

Checklist

Uh oh!

kartik-mem0 left a comment

Choose a reason for hiding this comment

Changes requested

Uh oh!

kartik-mem0 Jul 23, 2026

Choose a reason for hiding this comment

Uh oh!

kartik-mem0 Jul 24, 2026

Choose a reason for hiding this comment

Uh oh!

OfficialAbhinavSingh commented Jul 23, 2026

Uh oh!

kartik-mem0 left a comment

Choose a reason for hiding this comment

LGTM, ready for squash-merge

Uh oh!

OfficialAbhinavSingh commented Aug 9, 2026

Uh oh!

kartik-mem0 left a comment

Choose a reason for hiding this comment

Uh oh!

OfficialAbhinavSingh commented Sep 21, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants