fix(retriever): compute document query embeddings off the event loop - #1662
Conversation
DocumentHybridSearchRetriever called embedding_model.embed() directly from its coroutines. A local sentence-transformers model runs a synchronous forward pass, so that call stalls the event loop and every other request sharing the worker for its duration. Every sibling already offloads it: ChucksHybridSearchRetriever does in both vector_search and hybrid_search, and so do ConnectorService._combined_rrf and the multi-agent chat retrieval path. This retriever was the one that did not. Scope note: today these two lines are not reachable. vector_search and full_text_search have no callers, and hybrid_search's only caller always passes a precomputed query_embedding. So this is a consistency fix that keeps the next caller from silently reintroducing the stall -- not a measured speedup, and the PR should not be read as one.
|
@Yigtwxx 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 |
DocumentHybridSearchRetrievercallsembedding_model.embed()synchronously inside its coroutines. Every sibling in the codebase offloads that call; this one file does not.Description
A local sentence-transformers model runs a synchronous forward pass taking tens to hundreds of milliseconds. Calling it directly from a coroutine stalls the event loop — and therefore every other request on that worker — for its duration.
The correct pattern is already established in four places:
app/retriever/chunks_hybrid_search.py:93await asyncio.to_thread(embedding_model.embed, query_text)app/retriever/chunks_hybrid_search.py:241app/services/connector_service.py:187app/agents/chat/multi_agent_chat/shared/retrieval/hybrid_search.py:91app/retriever/documents_hybrid_search.pylines 91 and 225 were the two that were missed — the direct sibling ofchunks_hybrid_search.py, in the same two methods.Motivation and Context
Consistency with the pattern the rest of the retrieval layer already follows.
Please read this part before weighing the PR. I checked whether these lines are reachable today, and they are not:
vector_searchandfull_text_searchon this class have no callers anywhere in the repo.hybrid_searchhas exactly one caller,ConnectorService._combined_rrf, which always passes a precomputedquery_embedding— so theif query_embedding is Nonebranch does not execute.So this is not a measured latency improvement, and I would rather say that up front than let the diff imply one. The value is that the file stops disagreeing with its four siblings, and the next caller to use these entry points does not silently reintroduce the stall. If you would rather delete the two unused methods than fix them, that is a reasonable alternative and I am happy to send that instead.
Screenshots
Not applicable — backend only.
API Changes
Change Type
Testing Performed
uv run pytest -m unit: 2845 passed, same 7 pre-existing failures asdevon this machine (git-tree and PAT static tests, unrelated and failing before this branch).Two tests added. They use a fake embedding model that blocks a real thread on a
threading.Event, then assert the coroutine has not completed while it is blocked — so the assertion is that the loop kept running, not that a particular function was called. Stashing only the source change:Both hang until the 5s release timeout and then fail on
dev, which is the stall itself being observed.Checklist
What does not change
chunks_hybrid_search.pyandconnector_service.pyare untouched; they were already correct.High-level PR Summary
This PR fixes
DocumentHybridSearchRetrieverto compute query embeddings off the event loop usingasyncio.to_thread, matching the pattern already established in four sibling files (chunks_hybrid_search.py,connector_service.py, and chat retrieval). The fix prevents potential event loop blocking when embedding models perform synchronous forward passes that can take tens to hundreds of milliseconds. While the affected code paths (vector_searchandhybrid_searchwithout precomputed embeddings) are currently unreachable in production, this change ensures consistency across the codebase and prevents future regressions when these methods are used.⏱️ Estimated Review Time: 5-15 minutes
💡 Review Order Suggestion
surfsense_backend/app/retriever/documents_hybrid_search.pysurfsense_backend/tests/unit/retriever/test_query_embedding_offloading.py