Skip to content

Implement document expansion to improve RAG performance - #1

Merged
Morgan-Sell merged 8 commits into
mainfrom
feat/doc-expansion
Apr 25, 2026
Merged

Implement document expansion to improve RAG performance#1
Morgan-Sell merged 8 commits into
mainfrom
feat/doc-expansion

Conversation

@Morgan-Sell

Copy link
Copy Markdown
Owner

Summary

Adds V1 document expansion after fused retrieval and before reranking. This helps the pipeline recover useful same-document context when retrieval finds the right source document but misses the most answer-bearing chunk.

The feature uses a route planner instead of a single expansion path:

  • Neighbor expansion by doc_id + chunk_index, supported across all chunking strategies.
  • Same-section expansion when section_path metadata is available.
  • Anchor-section expansion for field-like queries, such as “home resort,” to pull sections like Summary.
  • Aggressive caps to avoid flooding rerank/generation context.
  • Debug and trace visibility for chunks added by document expansion.

Why

Some multi-hop or field-style questions failed because retrieval identified the correct document but selected a nearby or related chunk instead of the chunk containing the key fact. For example, an athlete profile’s “Career Progression” chunk might be retrieved, while the “Summary” chunk contains the home resort.

Document expansion gives the reranker and generator a better candidate pool without requiring a broader initial retrieval sweep.

Implementation Notes

  • Added get_by_doc_id() support to the vector store protocol and Chroma backend.
  • Added src/retrieval/document_expansion.py for planning and applying expansion requests.
  • Wired expansion into RAGPipeline after fused retrieval and before rerank.
  • Added document_expansion config defaults.
  • Added trace/debug output via document_expansion_added.
  • Added tests for planner behavior, Chroma reads, caps, and a Red Gerard-style regression case.

Test Plan

  • uv run ruff check src/retrieval/document_expansion.py src/vectorstore/protocol.py src/vectorstore/chroma_store.py src/pipeline/rag_pipeline.py src/pipeline/models.py src/pipeline/trace.py src/config/loader.py src/evaluation/config_merge.py tests/test_document_expansion.py tests/test_vectorstore.py tests/test_config_merge.py
  • uv run pytest tests/test_document_expansion.py tests/test_vectorstore.py tests/test_config_merge.py tests/test_pipeline.py
  • uv run pytest

All tests passed: 143 passed, 2 deselected.

@Morgan-Sell
Morgan-Sell merged commit 3bbb8ba into main Apr 25, 2026
1 check passed
@Morgan-Sell
Morgan-Sell deleted the feat/doc-expansion branch April 25, 2026 02:34
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