Implement document expansion to improve RAG performance

.github/workflows/ci.yml

@@ -19,7 +19,12 @@ jobs:
           version: "0.6.x"
 
       - name: Install dependencies
-        run: uv sync --group dev
+        run: uv sync --group dev --frozen
+
+      - name: Debug Ruff
+        run: |
+          uv run ruff --version
+          sed -n '1,16p' tests/test_gold_qa.py
 
       - name: Ruff
         run: uv run ruff check src tests

README.md

@@ -1,7 +1,6 @@
 # Trailhead
 Trailhead is a snow-sports RAG app that answers questions over a richly linked world of athletes, resorts, competitions, and results. It pairs grounded retrieval, cited answers, evaluation tooling, and a clean UI, with just enough mountain energy to make the whole thing more fun than your average AI demo.
 
-`Hierarchical retrieval` `Cited answers` `RAG sweeps` `Trace analysis` `Dockerized UI`
 
 <!-- Best screenshot placement: insert a Trailhead UI screenshot here, directly under the intro and before the feature sections.
 Example:

configs/default.yaml

@@ -52,6 +52,20 @@ query_expansion:
   fusion: max_score
   rrf_k: 60
 
+document_expansion:
+  enabled: true
+  modes: [neighbors, same_section, anchor_sections]
+  window: 1
+  max_seed_hits: 8
+  max_extra_chunks_per_doc: 3
+  max_total_extra_chunks: 8
+  expansion_score_penalty: 0.05
+  anchor_sections_by_entity_type:
+    athletes: [Summary]
+    resorts: [Summary]
+    circuits: [Summary]
+    competitions: [Summary]
+
 generation:
   # Phase 2.4: grounded answers. When disabled, the CLI prints retrieval rows
   # only and does not call an LLM for synthesis.

pyproject.toml

@@ -45,7 +45,7 @@ packages = [
 [dependency-groups]
 dev = [
     "pytest>=8.0",
-    "ruff>=0.8",
+    "ruff==0.15.10",
 ]
 
 [tool.pytest.ini_options]
@@ -61,3 +61,6 @@ src = ["src", "app", "tests"]
 
 [tool.ruff.lint]
 select = ["E", "F", "I", "W"]
+
+[tool.ruff.lint.isort]
+known-first-party = ["snow_sports_rag"]

src/config/loader.py

@@ -69,6 +69,21 @@
         "fusion": "max_score",
         "rrf_k": 60,
     },
+    "document_expansion": {
+        "enabled": True,
+        "modes": ["neighbors", "same_section", "anchor_sections"],
+        "window": 1,
+        "max_seed_hits": 8,
+        "max_extra_chunks_per_doc": 3,
+        "max_total_extra_chunks": 8,
+        "expansion_score_penalty": 0.05,
+        "anchor_sections_by_entity_type": {
+            "athletes": ["Summary"],
+            "resorts": ["Summary"],
+            "circuits": ["Summary"],
+            "competitions": ["Summary"],
+        },
+    },
     "generation": {
         "enabled": False,
         "backend": "openai",
@@ -119,6 +134,8 @@ class AppConfig:
         LLM provider and generation parameters.
     query_expansion : Mapping[str, Any]
         Phase 2.2 multi-query expansion and fusion options.
+    document_expansion : Mapping[str, Any]
+        Post-retrieval same-document expansion options.
     generation : Mapping[str, Any]
         Phase 2.4 grounded answer generation (backend, prompt, model ids).
     logging : Mapping[str, Any]
@@ -133,6 +150,7 @@ class AppConfig:
     rerank: Mapping[str, Any]
     llm: Mapping[str, Any]
     query_expansion: Mapping[str, Any]
+    document_expansion: Mapping[str, Any]
     generation: Mapping[str, Any]
     logging: Mapping[str, Any]
 
@@ -244,6 +262,7 @@ def load_config(
         rerank=merged["rerank"],
         llm=merged["llm"],
         query_expansion=merged["query_expansion"],
+        document_expansion=merged["document_expansion"],
         generation=merged["generation"],
         logging=merged["logging"],
     )

src/evaluation/config_merge.py

@@ -76,6 +76,7 @@ def merge_app_config_overrides(
             "rerank",
             "llm",
             "query_expansion",
+            "document_expansion",
             "generation",
             "logging",
         }
@@ -102,6 +103,7 @@ def merge(name: str) -> dict[str, Any]:
         rerank=merge("rerank"),
         llm=merge("llm"),
         query_expansion=merge("query_expansion"),
+        document_expansion=merge("document_expansion"),
         generation=merge("generation"),
         logging=merge("logging"),
     )

src/gradio_app/components.py

@@ -354,6 +354,8 @@ def _rows_from_hits(hits: list, limit: int = 50) -> list[dict]:
             _render_string_list(trace.l1_shortlist),
             "<h4>L2 candidates (pre-rerank)</h4>",
             _render_hit_table(_rows_from_hits(trace.l2_pre_rerank)),
+            "<h4>Document expansion added</h4>",
+            _render_hit_table(_rows_from_hits(trace.document_expansion_added)),
             "<h4>Reranked / final</h4>",
             _render_hit_table(_rows_from_hits(trace.reranked)),
             "<h4>Latency</h4>",

src/pipeline/models.py

@@ -94,6 +94,8 @@ class PipelineTrace:
         Union of L1 doc_ids across variants, in first-seen order.
     l2_pre_rerank : list of RetrievalHit
         Fused L2 candidates before the reranker is applied.
+    document_expansion_added : list of RetrievalHit
+        Extra same-document chunks added after fused retrieval.
     reranked : list of RetrievalHit
         Output of the reranker; identical to ``l2_pre_rerank[:top_k_out]``
         when reranking is disabled.
@@ -106,6 +108,7 @@ class PipelineTrace:
     variants: list[str] = field(default_factory=list)
     l1_shortlist: list[str] = field(default_factory=list)
     l2_pre_rerank: list[RetrievalHit] = field(default_factory=list)
+    document_expansion_added: list[RetrievalHit] = field(default_factory=list)
     reranked: list[RetrievalHit] = field(default_factory=list)
     latency: StageLatency = field(default_factory=StageLatency)
 

src/pipeline/rag_pipeline.py

@@ -8,6 +8,7 @@
 from ..generation import answer_generator_from_config
 from ..llm import llm_client_from_config
 from ..rerank import reranker_from_config
+from ..retrieval.document_expansion import expand_retrieval_hits
 from ..retrieval.fusion import (
     fuse_retrieval_hits_max_score,
     fuse_retrieval_hits_rrf,
@@ -441,6 +442,15 @@ def run(
         l2_pre_rerank = self._fuse(
             per_variant_hits, top_n_fused=preset_obj.top_n_pre_rerank
         )
+        document_expansion_added: list[RetrievalHit] = []
+        if bool(self._cfg.document_expansion.get("enabled", False)):
+            assert self._l2_store is not None
+            l2_pre_rerank, document_expansion_added, _ = expand_retrieval_hits(
+                l2_pre_rerank,
+                query=q,
+                store=self._l2_store,
+                config=self._cfg.document_expansion,
+            )
         retrieval_ms = (time.perf_counter() - t0) * 1000.0
 
         t0 = time.perf_counter()
@@ -484,6 +494,7 @@ def run(
             variants=list(variants),
             l1_shortlist=list(all_shortlists),
             l2_pre_rerank=list(l2_pre_rerank),
+            document_expansion_added=list(document_expansion_added),
             reranked=list(reranked),
             latency=latency,
         )

src/pipeline/trace.py

@@ -60,7 +60,7 @@ def compute_config_hash(cfg: Any) -> str:
 
     Only sections that materially change retrieval / generation behaviour are
     included (chunking, embedding, vector_store, retrieval, rerank,
-    query_expansion, generation). Secrets are stripped.
+    query_expansion, document_expansion, generation). Secrets are stripped.
 
     Parameters
     ----------
@@ -88,6 +88,7 @@ def _get(key: str) -> Any:
             "retrieval",
             "rerank",
             "query_expansion",
+            "document_expansion",
             "generation",
         )
     }
@@ -245,6 +246,9 @@ def log_query(
             "variants": list(tr.variants),
             "l1_shortlist": list(tr.l1_shortlist),
             "l2_pre_rerank": [_hit_to_dict(h) for h in tr.l2_pre_rerank],
+            "document_expansion_added": [
+                _hit_to_dict(h) for h in tr.document_expansion_added
+            ],
             "reranked": [_hit_to_dict(h) for h in tr.reranked],
             "final_sources": [
                 {