Skip to content

fix: make reranker propely optional - #1690

Open
Benebo7 wants to merge 2 commits into
MODSetter:devfrom
Benebo7:fix/reimpementing-reranker
Open

fix: make reranker propely optional#1690
Benebo7 wants to merge 2 commits into
MODSetter:devfrom
Benebo7:fix/reimpementing-reranker

Conversation

@Benebo7

@Benebo7 Benebo7 commented Aug 17, 2026

Copy link
Copy Markdown
Contributor

Summary

RERANKERS_ENABLED is documented in both .env.example files and config builds a live Reranker when it's set, but search_knowledge_base calls build_context(query, hits, registry) without the reranker argument. So rerank_hits always receives None and returns hits untouched. The switch exists and does nothing.

This wires it through. rank() is a blocking cross-encoder call, so the enabled path goes through asyncio.to_thread to keep it off the event loop. The disabled path stays exactly as it is today.

reranker = RerankerService.get_reranker_instance()
if reranker is not None:
    rendered = await asyncio.to_thread(
        build_context, cleaned_query, hits, registry, reranker=reranker
    )
else:
    rendered = build_context(cleaned_query, hits, registry)

get_reranker_instance() returns None whenever RERANKERS_ENABLED is off, so nothing changes for anyone who doesn't opt in: same call, same code path, no thread. The default stays FALSE, so cloud is unaffected. What changes is that self-hosted operators can now actually turn it on, and cloud operators can actually test and measure metrics if wanted.

Testing

Measured on the dev stack against a real indexed workspace.

  • Flag OFF: build_context costs ~2 ms.
  • Flag ON: Reranking adds 0.8 to 2.1 s per search depending on document count.

A typical turn in my logs spends 46 to 72 seconds, almost all of it in LLM calls, so reranking lands around 5% of turn time.

Open Questions (Not in this PR)

Wiring the reranker up surfaced two things that affect how it performs. Both are judgment calls that depend on knowing the cloud setup, so I've left them out and am raising them here instead.

  1. The default model: .env.example suggests ms-marco-MiniLM-L-12-v2 via flashrank. ms-marco-MiniLM-L-6-v2 scores the same on MS MARCO (74.30 vs 74.31 NDCG@10) and was consistently faster in my measurements: roughly 3x on small result sets, narrowing to 1.4x once documents get large enough that both truncate. The catch is that FlashRank doesn't ship L-6, so using it means switching RERANKERS_MODEL_TYPE to cross-encoder, which means PyTorch instead of ONNX.
  2. Concurrency: The agent fires searches in parallel. Since the cross-encoder releases the GIL and runs in native code, those genuinely compete for cores. Without a limit they all finish at the slowest one's pace. A semaphore around the reranking call gave each search predictable latency, and at higher concurrency it improved wall-clock time too. Whether that helps at all depends on whether reranking threads actually exceed available cores on your infrastructure. Sizing it also needs to account for worker count, since this optimization is targeting cloud, and I can't tell how the cloud deployment handles that.

Trade-off Summary

Decision Upside Cost
Switch default to L-6 1.4 to 3x faster, same benchmark quality Needs cross-encoder backend (PyTorch, not ONNX); +24 MB RSS; 1.4 s model load vs 0.2 s
Keep L-12 / flashrank ONNX runtime purpose-built for inference; faster startup Roughly half the throughput for no measurable quality gain
Add a concurrency limit Predictable per-search latency; better wall-clock under load Adds queueing when cores are free; sizing depends on cores and worker count
No limit (today) No queueing when capacity is free Parallel searches contend; all finish at the slowest one's pace

Happy to open follow-ups for either once you've had a look. You'd know better than I would whether the cloud hits the concurrency case at all.

Note: build_context only touches plain dataclasses, so no lazy load can cross into the worker thread.

High-level PR Summary

This PR fixes the reranker feature to actually work when enabled. Previously, the RERANKERS_ENABLED configuration flag existed but was never used because the reranker argument was not passed to build_context(), causing search results to never be reranked. The fix conditionally retrieves the reranker instance and, when enabled, wraps the build_context() call in asyncio.to_thread() to prevent blocking the event loop during the CPU-intensive cross-encoder operation. When the flag is disabled (the default), behavior remains identical to the current implementation.

⏱️ Estimated Review Time: 5-15 minutes

💡 Review Order Suggestion
Order File Path
1 surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/knowledge_base/tools/search_knowledge_base.py

Need help? Join our Discord

@vercel

vercel Bot commented Aug 17, 2026

Copy link
Copy Markdown

@Benebo7 is attempting to deploy a commit to the Rohan Verma's projects Team on Vercel.

A member of the Team first needs to authorize it.

@coderabbitai

coderabbitai Bot commented Aug 17, 2026

Copy link
Copy Markdown

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 91f8a226-d3ba-4e7f-8479-70a3888efda0

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

  • 🔍 Trigger review

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@MODSetter

Copy link
Copy Markdown
Owner

@AnishSarkar22 Can you review this?

@Yigtwxx Yigtwxx left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not a maintainer — take this as input rather than a gate.

I measured refs/pull/1690/merge rather than the branch, since that is the tree CI checks out.

The diagnosis holds independently. I hit the same gap from the other direction while comparing the legacy retriever against shared/retrieval, before reading this PR: build_context has exactly one production call site (search_knowledge_base.py:157) and it never passed reranker=, and RerankerService.get_reranker_instance() had no non-test caller at all. So RERANKERS_ENABLED really was a switch wired to nothing.

Two things I checked because they are the parts that could bite, both of which came out fine:

1. No AsyncSession crosses the thread boundary. The to_thread call sits inside async with shielded_async_session() as session, which is the shape that usually hides a lazy load. It does not here — I traced the whole callee set:

  • DocumentHit (shared/retrieval/models.py:36) is a plain dataclass; _group_into_documents fills title, document_type, metadata and the chunk contents eagerly, so nothing is a deferred ORM attribute.
  • to_renderable_document (shared/retrieval/adapter.py) only reads those fields.
  • render_search_context / render_document / source_label are pure string builders with no await and no session reference.

So build_context is genuinely CPU-only, and moving it to a worker thread is safe. Worth stating in the description, because "we moved a call that lives inside a DB session into a thread" is the first thing a reviewer will worry about.

2. The enabled path does not reload the model per search. config.reranker_instance is constructed once, at class-body evaluation in app/config/__init__.py:1036, so the weights load at import. get_reranker_instance() only wraps that instance in a new RerankerService, which is cheap. The 0.8–2.1 s you measured is inference, not loading — which matches what you reported and is the answer to the obvious "does the first search pay for the model?" question.

One suggestion and one nit:

  • Resolve the reranker once, at tool construction. get_reranker_instance() is called inside _impl, so every search re-reads config and allocates a wrapper. Since config.reranker_instance is fixed at import, create_search_knowledge_base_tool could resolve it alongside _space_id and _document_types and close over it. Same behaviour, one less allocation per search, and it makes the disabled path a plain closure check.

  • With that hoisted, the if reranker is not None / else pair collapses to one branch that either awaits to_thread or calls directly — as written the two arms duplicate the argument list, which is the kind of thing that drifts later.

On your two open questions: both look like maintainer calls to me and I would not fold either into this PR. For what it is worth, the concurrency one has a precedent in the tree — app/utils/document_converters.py:19 already serialises embedding-model access with an RLock for a related reason (HF fast tokenizers are not thread-safe), so a bounded-concurrency wrapper around the reranker would not be a new pattern here.

@Benebo7

Benebo7 commented Aug 21, 2026

Copy link
Copy Markdown
Contributor Author

Yo, thanks for reviewing my code once again @Yigtwxx , that's some good observations
Hope code is better now. I actually just added reranker to non reranking build context call parameters so both function calls can be similar, i guess it resolves your nit.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants