ref(kb): H04 move search lifecycle into handle (#511)#670
Conversation
Re-route the legacy-step compatibility facade's search_dense/search_sparse/ search_hybrid (+ async dense/sparse) to open a READ collection handle via the coordinator and delegate to handle.search_*, instead of the old _search_*_impl free functions. Lift the public search_dense() DocumentValidationError input validation to its boundary (the routed path no longer passes through the impl). Migrate the retrieval unit tests to the new handle seam (shared conftest fixtures) without weakening any behavioral assertion. Claude-Session: https://claude.ai/code/session_01LcfvwrLffsY8iGF2XgbhFu
…ai#511) - Delete search_engine.py (4 dead functions: search_dense_engine, _search_dense_engine_impl, search_dense_engine_async, _search_dense_engine_async_impl) - Delete _search_dense_impl/_search_dense_async_impl from search_dense.py; drop now-unused search_dense_engine import - Delete _search_sparse_impl/_search_sparse_async_impl/_substring_fallback/ _substring_fallback_async/_build_sparse_response from search_sparse.py - Delete _search_hybrid_impl from search_hybrid.py; drop private impl imports - Delete search_dense_engine/search_dense_engine_async facade methods from retrieval_compatibility.py; clean up now-unused imports Test changes: - Delete TestSearchDenseEngine + TestSearchDenseIntegration engine tests from test_search_dense.py (behavior covered by handle tests) - Port Issue xorbitsai#72 exact collection-filter equality assertions to test_collection_handle_search.py (new tests verify FilterCondition( field=collection, operator=EQ) reaches store) - Remove all facade.search_dense_engine tests from test_retrieval_helper_compatibility.py; keep surface-import + LLM-format tests - Update test_multitenancy.py to call search_dense() instead of search_dense_engine() (response object instead of tuple) - Update test_search_sparse.py substring-fallback test to call via handle method Claude-Session: https://claude.ai/code/session_01LcfvwrLffsY8iGF2XgbhFu
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the RAG search data-plane by moving dense, sparse, and hybrid search implementations from free functions into the LanceDBCollectionHandle class, with public search functions now delegating to the KB coordinator facade. The review feedback identifies several critical issues: an unnecessary database fetch in _substring_fallback_async due to a missing outer loop break, potential crashes in both sync and async substring fallbacks when non-dictionary filters are passed, and potential KeyErrors in _dense_engine when accessing optional fields directly instead of using .get().
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
…tsai#511) Address defensive-hardening review findings (verbatim pre-existing behavior preserved from the refactored search helpers): - guard substring-fallback filter handling with isinstance(dict) checks so non-dict filters no longer raise AttributeError/TypeError - add the missing outer-loop break in the async substring fallback so it stops fetching batches once top_k is reached, matching the sync path - read parse_hash/created_at via row.get() in the sync dense engine to match the async engine and avoid KeyError on nullable columns Claude-Session: https://claude.ai/code/session_01LcfvwrLffsY8iGF2XgbhFu
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request refactors the RAG search data-plane by consolidating dense, sparse, and hybrid search logic (both sync and async) directly into the LanceDBCollectionHandle and updating legacy compatibility layers to delegate to it. Feedback on these changes highlights several critical improvements: supporting non-dictionary filter types in the dense search engines for consistency with sparse search, correctly parsing legacy filter operators in the substring fallback method to prevent zero-result failures, wrapping the optional PyArrow dependency in a try-except block to avoid import-time errors, and defensively clamping distances to prevent potential division-by-zero errors or invalid scores.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
…itsai#511) Low-risk defensive findings from PR review: - wrap the pyarrow import in try/except so a missing optional dependency degrades to the existing `pa is None` guard instead of failing at import - clamp distance with max(0.0, distance) in the dense score formula so a negative distance (e.g. dot-product metrics) cannot divide by zero or yield an out-of-range score Claude-Session: https://claude.ai/code/session_01LcfvwrLffsY8iGF2XgbhFu
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request refactors the RAG search data-plane by centralizing the core search execution logic (dense, sparse, and hybrid) inside the LanceDBCollectionHandle class, while updating the public retrieval functions to perform input validation and delegate to this handle. The code review feedback highlights several critical issues with filter handling in the new implementation: both dense search engines and the substring fallback methods silently ignore non-dict FilterExpression inputs, which could lead to data leakage across collections or users. Additionally, the reviewer suggests importing DocumentValidationError in the sparse and hybrid search modules to ensure proper input validation at the public boundary.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
qinxuye
left a comment
qinxuye
left a comment
There was a problem hiding this comment.
I found one async validation regression that should be fixed before approval.
…sai#511) The H04 refactor lifted dense input validation from the deleted _search_dense_impl to the public search_dense(), but the async path (search_dense_async) was left delegating straight to the facade, dropping the DocumentValidationError/VectorValidationError contract the original _search_dense_async_impl enforced. Extract the validation into a shared _validate_dense_inputs helper and apply it on both the sync and async public boundaries, with async invalid-input test coverage. Reported in PR review by @qinxuye. Claude-Session: https://claude.ai/code/session_01LcfvwrLffsY8iGF2XgbhFu
3 participants
Summary
Part of epic #494, Phase 2 ("KBCollectionHandle Replacement") — follows the merged H01/H02/H03 (#648/#656/#658).
Moves dense/sparse/hybrid search execution mechanics + response-model construction into
LanceDBCollectionHandle, re-routes the legacy-step compatibility facade through the coordinator-opened handle, and retires the now-dead engine layer. Behavior-preserving: public functions, import paths, rerank, and pipeline orchestration are unchanged.Closes #511.
What changed
kb/collection_handle.py): addedsearch_dense/search_sparse/search_hybrid(+ async dense/sparse and a new async hybrid) that buildDenseSearchResponse/SparseSearchResponse/HybridSearchResponsedirectly. A coarsecapabilities.supports_searchguard sits before the try/except and degrades to a failed response with aSEARCH_NOT_SUPPORTEDwarning (never raises). A shared synchronous_fuse_hybridhelper backs both hybrid paths.kb/legacy_step_compatibility.py):search_*(+ async) open a READ handle (hide_missing=True) and delegate tohandle.search_*. Sync viaopen_collection_sync, async viaawait coordinator.open_collection(...)(no separate event loop)._open_collection_handlegained anaccess_modeparameter (defaults to WRITE; parse/chunk callers unchanged)._search_*_impl, theretrieval/search_engine.pydense engine (file removed), andretrieval_compatibilityengine wrappers. Kept_rrf_fusion/_linear_fusion(shared with the pipeline),format_search_results_for_llm, and all publicsearch_*entry points. The publicsearch_dense()input validation (DocumentValidationError) was lifted to its boundary since the impl is gone.Behavior preserved (verified byte-identical)
Dense score
1/(1+distance); sparse TF-IDF normscore/(1+score); substring-fallback score1.0; warning codesREADONLY_MODE/FTS_INDEX_MISSING/FTS_FALLBACK/FTS_SEARCH_FAILED/DENSE_SEARCH_FAILED; RRF + linear fusion; index-status mapping; hybridtop_k*2over-fetch and score/rank attachment; collection isolation + user/admin filtering. Issue #72 collection-isolation coverage was ported to the handle tests at stronger strength (exact collection-filter equality).Testing
tests/core/tools/core/RAG_tools/— 1556 passed, 1 skipped, 1 pre-existing unrelated failure (version_management/test_promote_version_main.py::test_default_lancedb_dir_when_missing_env, an env-path assertion that also fails on the base). New handle search unit tests + a facade routing test added; the retrieval suite was migrated to exercise the handle seam without weakening assertions.ruffandmypy src/xagentclean on changed files.https://claude.ai/code/session_01LcfvwrLffsY8iGF2XgbhFu