Skip to content

feat(retrieval): scope lexical chunk search to node sets - #4675

Open
chinmayv095 wants to merge 1 commit into
topoteretes:devfrom
chinmayv095:feat/lexical-nodeset-filter
Open

feat(retrieval): scope lexical chunk search to node sets#4675
chinmayv095 wants to merge 1 commit into
topoteretes:devfrom
chinmayv095:feat/lexical-nodeset-filter

Conversation

@chinmayv095

Copy link
Copy Markdown
Contributor

Description

SearchType.CHUNKS takes node_name and node_name_filter_operator and hands them to the vector search, which filters on the belongs_to_set tags carried in each chunk payload. SearchType.CHUNKS_LEXICAL takes neither. The factory builds BM25ChunksRetriever with top_k alone, and LexicalRetriever.initialize() loads every DocumentChunk in the graph:

SearchType.CHUNKS_LEXICAL: (BM25ChunksRetriever, {"top_k": top_k}),

So a caller who scopes CHUNKS to one node set and then switches to CHUNKS_LEXICAL gets 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_model treats belongs_to_set as 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). Since LexicalRetriever is reading those same node payloads out of get_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 implementing get_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. BM25ChunksRetriever derives 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() raises NoDataError("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 no DocumentChunk nodes 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_entities also writes this property and reads it back through shapes that include mappings and DataPoints, so nodeset_tags reduces 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_tags into 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.

JaccardChunksRetriever gets the same two parameters, since it is the other LexicalRetriever subclass and would otherwise be the next thing to silently ignore a scope.

Acceptance Criteria

  • SearchType.CHUNKS_LEXICAL accepts node_name and node_name_filter_operator and applies them with the same meaning SearchType.CHUNKS gives them: OR keeps a chunk tagged with any requested set, AND only one tagged with all of them.
  • Unscoped lexical search is unchanged.
  • BM25 corpus statistics reflect only the scoped chunks.
  • An empty node set returns no results; a genuinely empty corpus still raises NoDataError.

9 new tests, in the existing test_bm25_retriever.py and test_get_search_type_retriever_instance.py rather than in new files:

cognee/tests/unit/modules/retrieval/  and  cognee/tests/unit/modules/search/
578 passed, 10 failed, 1 skipped

The 10 failures are identical, by name, on upstream/dev with these files reverted (569 passed there, so all 9 new tests are additive and nothing regressed). 8 of the 9 fail against dev; the 9th is test_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 --check is clean on all six files. ruff check reports no new rule violations against the base beyond the pre-existing Optional[...] / List[...] style already used throughout these files, which I matched rather than modernized so the new parameters read like their neighbours. ty check is unaffected: cognee/modules/retrieval and cognee/modules/search are not in [tool.ty.src] include, and the 9 diagnostics it reports are the same ones as on dev, none in these files.

Type of Change

  • Bug fix (non-breaking change that fixes an issue)
  • New feature (non-breaking change that adds functionality)
  • Code refactoring
  • Other (please specify):

Screenshots

gh cannot 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

  • I have tested my changes thoroughly before submitting this PR (See CONTRIBUTING.md)
  • This PR contains minimal changes necessary to address the issue/feature
  • My code follows the project's coding standards and style guidelines
  • I have added tests that prove my fix is effective or that my feature works
  • I have added necessary documentation (if applicable)
  • All new and existing tests pass
  • I have searched existing PRs to ensure this change hasn't been submitted already
  • I have linked any relevant issues in the description
  • My commits have clear and descriptive messages

Note on the existing-PR check: #3894 also touches lexical_retriever.py, but only get_context_from_objects, and it adds test_lexical_retriever.py. This PR touches __init__ and initialize() and puts its tests in the existing test_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.

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.
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.

1 participant