fix: make reranker propely optional - #1690
Conversation
|
@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. |
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
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. Comment |
|
@AnishSarkar22 Can you review this? |
Yigtwxx
left a comment
There was a problem hiding this comment.
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_documentsfillstitle,document_type,metadataand 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_labelare pure string builders with noawaitand 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-readsconfigand allocates a wrapper. Sinceconfig.reranker_instanceis fixed at import,create_search_knowledge_base_toolcould resolve it alongside_space_idand_document_typesand 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/elsepair collapses to one branch that either awaitsto_threador 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.
|
Yo, thanks for reviewing my code once again @Yigtwxx , that's some good observations |
Summary
RERANKERS_ENABLEDis documented in both.env.examplefiles andconfigbuilds a liveRerankerwhen it's set, butsearch_knowledge_basecallsbuild_context(query, hits, registry)without thererankerargument. Sorerank_hitsalways receivesNoneand 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 throughasyncio.to_threadto keep it off the event loop. The disabled path stays exactly as it is today.get_reranker_instance()returnsNonewheneverRERANKERS_ENABLEDis 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.
build_contextcosts ~2 ms.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.
.env.examplesuggestsms-marco-MiniLM-L-12-v2via flashrank.ms-marco-MiniLM-L-6-v2scores 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 switchingRERANKERS_MODEL_TYPEto cross-encoder, which means PyTorch instead of ONNX.Trade-off Summary
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_contextonly 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_ENABLEDconfiguration flag existed but was never used because thererankerargument was not passed tobuild_context(), causing search results to never be reranked. The fix conditionally retrieves the reranker instance and, when enabled, wraps thebuild_context()call inasyncio.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
surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/knowledge_base/tools/search_knowledge_base.py