feat(retrieval): scope lexical chunk search to node sets - #4675
Open
chinmayv095 wants to merge 1 commit into
Open
feat(retrieval): scope lexical chunk search to node sets#4675chinmayv095 wants to merge 1 commit into
chinmayv095 wants to merge 1 commit into
Conversation
SearchType.CHUNKS accepts node_name and node_name_filter_operator and passes them to the vector search, which filters on the belongs_to_set tags carried in each chunk payload. SearchType.CHUNKS_LEXICAL accepts nothing: the factory builds BM25ChunksRetriever with top_k alone, and LexicalRetriever loads every DocumentChunk in the graph. A caller who scopes CHUNKS to one node set and then switches to CHUNKS_LEXICAL gets the whole corpus back, with no indication the scope was dropped. Give the lexical retrievers the same two parameters and apply them while the corpus is being loaded. The tags are already on the node properties: graph serialization keeps belongs_to_set as a property, reduced to NodeSet names, specifically so node sets can be filtered on, so the filter needs no adapter support and no second graph query. OR keeps a chunk tagged with any requested set and AND one tagged with all of them, matching what the vector adapters do with the same field. Filtering at load time rather than at scoring time also keeps BM25's corpus statistics honest: IDF and average chunk length are derived from what was loaded, so a scoped search now ranks against the corpus it can actually return rather than against the whole graph. An empty node set returns no results instead of raising NoDataError, which claims the system holds no data at all. That matches the vector chunk search, which returns nothing in the same situation. A genuinely empty corpus still raises.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
SearchType.CHUNKStakesnode_nameandnode_name_filter_operatorand hands them to the vector search, which filters on thebelongs_to_settags carried in each chunk payload.SearchType.CHUNKS_LEXICALtakes neither. The factory buildsBM25ChunksRetrieverwithtop_kalone, andLexicalRetriever.initialize()loads everyDocumentChunkin the graph:So a caller who scopes
CHUNKSto one node set and then switches toCHUNKS_LEXICALgets the whole corpus back. The scope is not rejected and not warned about, it is just dropped, which is the failure mode you notice last.I went looking for where the tags live before deciding how to filter, and they turn out to be on the node already.
get_graph_from_modeltreatsbelongs_to_setas the one DataPoint-valued field that is not excluded from node properties, and reduces it to a list of NodeSet names, with a comment saying it does this so node sets can be filtered on. That is what the vector adapters read (payload -> 'belongs_to_set' ?| ARRAY[...]in PGVector). SinceLexicalRetrieveris reading those same node payloads out ofget_filtered_graph_data, the filter is a local check on data already in hand: no adapter support to add, no second graph query, and it works on every adapter where lexical search works today rather than only on those implementingget_nodeset_subgraph.Three decisions worth flagging, since none of them are forced by the code:
The filter runs while the corpus is loading, not while it is being scored. That is not only about skipping tokenization.
BM25ChunksRetrieverderives IDF and average chunk length from whatever the parent loaded, so filtering later would leave BM25 ranking a scoped result set against statistics computed over the entire graph. A term that is common inside the requested node set but rare outside it would keep an IDF that describes a corpus the search cannot return. Filtering first makes the statistics describe the corpus actually being ranked.An empty node set returns nothing instead of raising.
initialize()raisesNoDataError("No valid chunks loaded during initialization.")when it loads zero chunks. With scoping added, the same path is reached when the system is full of data and the requested node set is simply empty, and telling that caller there is no data in the system would be wrong. The vector chunk search returns an empty result there, so this now does too. A genuinely empty corpus, with no node-set filter or with noDocumentChunknodes at all, still raises exactly as before, and there is a test pinning each half.The tag normalizer tolerates more than the serializer emits. Serialization always produces plain names, but
consolidate_entitiesalso writes this property and reads it back through shapes that include mappings and DataPoints, sonodeset_tagsreduces all three to the same string key rather than assuming the happy path. I deliberately did not move that module's private_belongs_to_set_tagsinto shared code as part of this PR: it is a task-layer helper, importing it into retrieval would be the wrong direction, and unifying them properly is a refactor that does not belong in a feature change. Happy to do it as a follow-up if you would rather have one implementation.JaccardChunksRetrievergets the same two parameters, since it is the otherLexicalRetrieversubclass and would otherwise be the next thing to silently ignore a scope.Acceptance Criteria
SearchType.CHUNKS_LEXICALacceptsnode_nameandnode_name_filter_operatorand applies them with the same meaningSearchType.CHUNKSgives them:ORkeeps a chunk tagged with any requested set,ANDonly one tagged with all of them.NoDataError.9 new tests, in the existing
test_bm25_retriever.pyandtest_get_search_type_retriever_instance.pyrather than in new files:The 10 failures are identical, by name, on
upstream/devwith these files reverted (569 passed there, so all 9 new tests are additive and nothing regressed). 8 of the 9 fail againstdev; the 9th istest_no_node_name_searches_every_chunk, which passes both before and after on purpose, because its job is to pin that unscoped search did not change.ruff format --checkis clean on all six files.ruff checkreports no new rule violations against the base beyond the pre-existingOptional[...]/List[...]style already used throughout these files, which I matched rather than modernized so the new parameters read like their neighbours.ty checkis unaffected:cognee/modules/retrievalandcognee/modules/searchare not in[tool.ty.src] include, and the 9 diagnostics it reports are the same ones as ondev, none in these files.Type of Change
Screenshots
ghcannot attach images, so the run is pasted above rather than shown. Happy to add a screenshot if you want it on the PR.Pre-submission Checklist
CONTRIBUTING.md)Note on the existing-PR check: #3894 also touches
lexical_retriever.py, but onlyget_context_from_objects, and it addstest_lexical_retriever.py. This PR touches__init__andinitialize()and puts its tests in the existingtest_bm25_retriever.py, so the two do not overlap textually. No open PR references node-set filtering for lexical search, and there is no issue open for it that I could find.DCO Affirmation
I affirm that all code in every commit of this pull request conforms to the terms of the Topoteretes Developer Certificate of Origin.