Skip to content

fix(rag): count identical chunks once in a search and cite every copy - #7131

Merged
kovtcharov-amd merged 2 commits into
amd:mainfrom
ityashag:fix/rag-collapse-duplicate-chunks
Nov 10, 2026
Merged

kovtcharov-amd merged 2 commits into
amd:mainfrom
ityashag:fix/rag-collapse-duplicate-chunks

Conversation

@ityashag

Copy link
Copy Markdown
Contributor

Summary

A document indexed twice under different names no longer fills the model's context with copies of itself. Identical chunks now count once in a search, so the model gets k different chunks. Each collapsed chunk lists every file that holds it, and the agent cites all of them instead of whichever copy was indexed first.

Why

The exact_duplicate hard case in #4733 found duplicates taking 40% of the top 5 chunks given to the model. As the bot noted on the issue, collapsing after FAISS's top-k cut would only trade crowding for a short context, so the search fetches more candidates until it has k distinct ones.

Linked issue

Closes #4740

Changes

  • The SDK search collapses identical chunks before cutting to k. It asks FAISS for more results until it has k distinct chunks or has seen the whole index, and records where each chunk's copies sit. chunk_metadata gains source_files / source_paths (every file holding the text), and query_metadata["source_files"] includes them. source_file / source_path stay as they were.
  • query_documents cites every file holding the text: source_file is the first, and also_in lists the rest. It resolves positions in one pass instead of rag.chunks.index() per chunk.
  • query_specific_file looks a chunk up only in the requested file, so a twin in another document can't supply the page number.
  • The RAG SDK guide says identical chunks count once, and fixes its example's meta['source'], a key that doesn't exist.

Test plan

Run in a python:3.12 container:

  • pytest tests/unit/rag/test_duplicate_chunks.py: 7 passed. All 7 fail on main (a copy takes a place, a missing source, a citation from the wrong file).
  • pytest tests/unit/rag tests/test_rag.py tests/unit/test_chat_rag_budget.py tests/unit/chat/ui/test_sse_handler.py: 1,341 passed, 35 skipped.
  • Full pytest tests/unit: 38,405 passed, 584 skipped, 15 failed, of 38,908 collected, with no crashed workers. The 15 fail the same way on main in this container: 14 shell-guardrail and skill-grant tests refuse gh and npm because neither is installed there, and one expects a folder it can't create, but the container runs as root. 44 tests have :: inside their parameters, so they were rerun by file to account for them.
  • python util/lint.py --all: passes. Pylint crashed in its parallel workers, so it was rerun alone with --jobs 1, which exits 0.

The tool result the agent reads gains also_in, so rag_quality is the eval category to watch on the next RC.

Evidence

The new tests index one policy under policy.txt and policy (copy).txt beside three other documents, in a real FAISS index:

# main
AssertionError: a copy took one of the 3 places
AssertionError: assert '.../policy.txt' == '.../policy (copy).txt'   # page lookup used the other file

# this PR
7 passed
  • Agent exposed in the Agent UI: N/A, no visible change.
  • MCP tools / servers: N/A
  • CLI: N/A
  • HTTP API / REST: N/A

Checklist

  • I have linked a GitHub issue above (Closes #N / Fixes #N / Refs #N).
  • I have described why this change is being made, not just what changed.
  • I have run linting and tests locally (python util/lint.py --all, pytest tests/unit/).
  • I have attached real-world evidence matched to the surface I changed (see Evidence above), or marked each surface N/A.
  • I have updated documentation if user-visible behavior changed (see CONTRIBUTING.md).

A document indexed twice under different names filled the model's top-k
with copies of itself, and the agent cited whichever copy was indexed first.

The SDK search fetches more candidates until it has k distinct chunks and
records every file holding each one (chunk_metadata source_files and
source_paths). query_documents names the other files in also_in, and
query_specific_file resolves a chunk inside the requested file, so a twin
elsewhere cannot supply its page.

Closes amd#4740
@github-actions github-actions Bot added documentation Documentation changes rag RAG system changes tests Test changes performance Performance-critical changes agents labels Nov 10, 2026
@github-actions

Copy link
Copy Markdown
Contributor

Verdict: Request changes

This stops a document indexed twice under different names from filling the model's context with copies of itself: identical chunks now count once, the search fetches more until it has k different ones, and the agent cites every file holding the text. The approach is sound and well tested, but it opens one gap.

  • A copy can now vanish when one of its files is left out of a search. When a file is stale and fails to re-index, the search tool drops that file's chunks. Before this PR, the twin from the healthy copy still came back on its own. Now only one copy of each chunk survives the search, so if the surviving copy belongs to the left-out file, the text is dropped even though a good file still holds it. Keep the chunk when any of its copies belongs to a file that is being served, and add a test for that case.

Real-world evidence

N/A: no evidence bundle was produced for this run. The verdict rests on static review and the author's test output: 7 new tests that fail on main and pass here, plus the wider RAG suites. The agent's tool result gains an also_in field, so watch rag_quality on the next RC.

🔍 Technical details

🟡 _search_without drops a chunk whose kept copy is left out (src/gaia/agents/tools/rag_tools.py:451-465)

_search_chunks now keeps one index per distinct text and moves the others into snapshot["duplicate_indices"]. _search_without still checks only retrieved_indices[i] against left_out. Copies have identical embeddings, so FAISS breaks the tie however it likes. If it picks the left-out file's copy, the served twin goes with it. Before this PR both copies were in the top k, so the served one survived the filter. query_documents uses _chunk_positions(rag, left_out) afterwards and would correctly cite the served copy, but only if the chunk gets that far.

    snapshot = rag._search_chunks(query)
    owner = snapshot["chunk_to_file"]
    copies = snapshot.get("duplicate_indices") or [[] for _ in snapshot["retrieved_indices"]]
    kept = [
        (chunk, score)
        for idx, chunk, score, twins in zip(
            snapshot["retrieved_indices"],
            snapshot["retrieved_chunks"],
            snapshot["scores"],
            copies,
        )
        if any(owner.get(i) not in left_out for i in [idx, *twins])
    ]
    return [chunk for chunk, _ in kept], [score for _, score in kept]

For the test, mark one of the two policy files as left out (for example by patching _bring_current to return it), then check that the policy chunk is still returned and cites only the served file.

Checked, no issue

  • query_specific_file resolves positions from file_to_chunk_indices[str(target_file)], the same key the method already uses at line 1712.
  • The search loop always ends: fetch doubles up to len(chunks_snapshot) and stops there.
  • The retrieval eval harness (gaia/eval/retrieval/harness.py) reads retrieved_indices, so its numbers will now count duplicates once. That is intended, but worth knowing before comparing retrieval scores from before and after this change.
  • Doc drift: no other doc mentions meta['source'] or the chunk_metadata keys this PR changes.

Strengths

  • The PR fetches more candidates until it has k distinct chunks rather than collapsing after the top-k cut. That avoids trading crowding for a shorter context, the trap called out on the issue.
  • Replacing rag.chunks.index() per chunk with one pass over the index removes an O(N·k) scan.
  • The tests use a real FAISS index and real indexing, and each one fails on main for the reason it names.

Identical chunks now collapse to one in a search, so when the copy the
search kept belonged to a file left out of the search, the text vanished
even though another file still being served holds it. A chunk now stays
when any copy of its text sits in a served file.
@kovtcharov-amd
kovtcharov-amd added this pull request to the merge queue Nov 10, 2026
Merged via the queue into amd:main with commit c7d0a2f Nov 10, 2026
45 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

agents documentation Documentation changes performance Performance-critical changes rag RAG system changes tests Test changes

Projects

None yet

Development

Successfully merging this pull request may close these issues.

fix(rag): duplicate documents crowd the top-k chunks given to the model

2 participants