[FEATURE] Reduce entity-linking noise by keeping only the best-matching node per mention

[oussamahansal]

Description

The entity linker unions every top-k candidate for every extracted mention into the retrieval seed set, pulling in low-relevance nodes: three exact-match mentions produce eight seeds on the demo KG.

Add an opt-in single_best_match flag (default off) on ByoKGQueryEngine that keeps only the best-matching node per extracted mention. Off by default, so the seed set is unchanged; draft answers are never pruned.

This needs per-mention candidate grouping, which the matcher path destroys by flattening and globally sorting hits across all mentions. Add EntityLinker.link_grouped, which queries the index once per mention so attribution is preserved; link() is unchanged

Changes

• ByoKGQueryEngine(single_best_match=False) — new opt-in flag. When on, each extracted mention contributes only its best-matching node.
• EntityLinker.link_grouped(...) — new method returning candidates grouped per mention.

Problem

Related issue (if any): #

Testing

• Unit tests added/updated
• Integration tests added (as appropriate)
• Existing tests pass (pytest)
• Tested manually (describe below)

Checklist

• Code follows existing style and conventions
• License headers present on new files
• Documentation updated (if applicable)
• No breaking changes (or clearly documented)

---

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

[github-actions]

BYOKG-RAG Coverage Report: The coverage is at 94.86% (target: 80%). Download the HTML report here.

[mykola-pereyma]

Nice, well-contained change — the single_best_match flag is opt-in and off by default, the existing seed set is provably unchanged when off (the test asserts link_grouped is never called and seeds are identical), and draft answers are never pruned. I also verified the parts the tests mock against the real index source: index.query(input, topk=1, id_selector=None) matches the positional call, and document_id is a scalar string in dense/graph_store/fuzzy, so the seed-set shape is correct. Approving.

Three recommendations (none blocking, but worth addressing — see inline comments for two of them):

1. link_grouped (index.query) and link() (index.match) aren't equivalent for the fuzzy index. For the embedding indexes (DenseFaiss, GraphStore), match() is a batched loop of the same per-input search as query(), so per-mention results agree. But FuzzyStringIndex.match() applies a max_len_difference=4 short-string filter (and a global re-sort) that query() does not — so with a fuzzy matcher, single_best_match can select a "best" candidate the normal union path would have filtered out. Consider routing link_grouped through the matcher's own per-query logic (or applying the same length filter), or documenting that single_best_match assumes an embedding-backed matcher. A quick check: on a FuzzyStringIndex, compare link() vs link_grouped per mention and confirm the top candidate agrees.

2. The description says "link() is unchanged," but link()'s signature dropped the id_selector kwarg. That kwarg was dead (it was never forwarded — the EntityMatcher path uses match(), which has no id_selector), so behavior is unchanged, but it's still a public signature change: a caller passing link(..., id_selector=[...]) would now TypeError. Please correct the wording and note the removed kwarg.

3. Please complete the PR checklist — unit tests were added, but every box is currently unticked.

Thanks for keeping this opt-in and leaving the default path untouched — that's the right call for a seed-set change.

[github-actions]

BYOKG-RAG Coverage Report: The coverage is at 94.84% (target: 80%). Download the HTML report here.

[github-actions]

BYOKG-RAG Coverage Report: The coverage is at 95.06% (target: 80%). Download the HTML report here.

[noel-improv]

Nice work, just a few callouts; link_grouped duplicates link()'s body, and the two have drifted: id_selector only appears on one, and they raise at different points. Adding a group_by_mention flag to link() would drop the hasattr check in the engine. Declared on the Linker ABC with a default that loops over link([m]), it would leave third-party linkers unchanged. The guard doesn't hold anyway, since entity_linker is public and reassigning it after construction still triggers the mid-query AttributeError it's meant to prevent.

Calling retrieve() per mention keeps inputs at length 1, so a dense index runs N embedding round trips instead of one batched call. topk isn't forwarded either, so each search runs at the linker's default width of 3 and keeps one hit. Since parse_response doesn't dedupe and LLMs repeat entities, a repeated mention now costs a second round trip instead of collapsing into one batch.

single_best_match also prunes explored_entities, which feeds the path retriever, so path coverage narrows too. The docstring and configuration.mdx mention only the triplet seeds.

Two smaller nits. The global-sort justification in the docstring applies only to FuzzyStringIndex. Both dense indexes already return hits grouped per input, per the same wording at graph-retrievers.mdx:63. single_best_match is also missing from the full-initialization example in query-engine.mdx, so that list disagrees with configuration.mdx:30.

[github-actions]

BYOKG-RAG Coverage Report: The coverage is at 95.04% (target: 80%). Download the HTML report here.