From b628dd0981c00c22e91c4a64f3c187914eb690b5 Mon Sep 17 00:00:00 2001
From: Morgan-Sell <morganpsell@gmail.com>
Date: Fri, 24 Apr 2026 18:39:46 -0700
Subject: [PATCH 1/8] Implement document expansion to improve RAG performance

---
 README.md                           |   1 -
 configs/default.yaml                |  14 ++
 img/trailhead_ui.png                | Bin 0 -> 340 bytes
 src/config/loader.py                |  19 ++
 src/evaluation/config_merge.py      |   2 +
 src/gradio_app/components.py        |   2 +
 src/pipeline/models.py              |   3 +
 src/pipeline/rag_pipeline.py        |  11 ++
 src/pipeline/trace.py               |   6 +-
 src/retrieval/document_expansion.py | 262 ++++++++++++++++++++++++++++
 src/vectorstore/chroma_store.py     |  32 ++++
 src/vectorstore/protocol.py         |  18 +-
 tests/test_config_merge.py          |   1 +
 tests/test_document_expansion.py    | 219 +++++++++++++++++++++++
 tests/test_vectorstore.py           |  37 ++++
 15 files changed, 624 insertions(+), 3 deletions(-)
 create mode 100644 img/trailhead_ui.png
 create mode 100644 src/retrieval/document_expansion.py
 create mode 100644 tests/test_document_expansion.py

diff --git a/README.md b/README.md
index bead92a..2d137de 100644
--- a/README.md
+++ b/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:
diff --git a/configs/default.yaml b/configs/default.yaml
index 257d1f6..70cdd34 100644
--- a/configs/default.yaml
+++ b/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.
diff --git a/img/trailhead_ui.png b/img/trailhead_ui.png
new file mode 100644
index 0000000000000000000000000000000000000000..4ed834d99313629eaeb3199c37cca35e00e0da42
GIT binary patch
literal 340
zcmeAS@N?(olHy`uVBq!ia0y~yV2WU1V4T3g3>0yF{caVIVh->Ladj{7Y5H~P5|GPS
z666=m;PC858j$nf)5S5Q;?~=PjJ!bJp#|UlbNSgtfXopAJ|wNSWngsg%ijs|uBWS?
J%Q~loCIHs1A6ozb

literal 0
HcmV?d00001

diff --git a/src/config/loader.py b/src/config/loader.py
index 080141c..ba45624 100644
--- a/src/config/loader.py
+++ b/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"],
     )
diff --git a/src/evaluation/config_merge.py b/src/evaluation/config_merge.py
index c137e2d..662aa0a 100644
--- a/src/evaluation/config_merge.py
+++ b/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"),
     )
diff --git a/src/gradio_app/components.py b/src/gradio_app/components.py
index 4bc291d..66bd37e 100644
--- a/src/gradio_app/components.py
+++ b/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>",
diff --git a/src/pipeline/models.py b/src/pipeline/models.py
index 926e56b..af64fab 100644
--- a/src/pipeline/models.py
+++ b/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)
 
diff --git a/src/pipeline/rag_pipeline.py b/src/pipeline/rag_pipeline.py
index 8045ceb..7f1e0df 100644
--- a/src/pipeline/rag_pipeline.py
+++ b/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,
         )
diff --git a/src/pipeline/trace.py b/src/pipeline/trace.py
index 2adb384..5b60731 100644
--- a/src/pipeline/trace.py
+++ b/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": [
                 {
diff --git a/src/retrieval/document_expansion.py b/src/retrieval/document_expansion.py
new file mode 100644
index 0000000..a308fc3
--- /dev/null
+++ b/src/retrieval/document_expansion.py
@@ -0,0 +1,262 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Literal, Mapping, Protocol
+
+from ..vectorstore.models import VectorQueryHit
+from .models import RetrievalHit
+from .scoring import chroma_cosine_distance_to_similarity
+
+__all__ = [
+    "ExpansionRequest",
+    "SupportsGetByDocId",
+    "expand_retrieval_hits",
+    "plan_expansion_requests",
+]
+
+__doc__ = """Post-retrieval document expansion helpers."""
+
+
+ExpansionMode = Literal["neighbors", "same_section", "anchor_sections"]
+
+
+class SupportsGetByDocId(Protocol):
+    """Vector-store read surface needed for document expansion."""
+
+    def get_by_doc_id(self, doc_id: str) -> list[VectorQueryHit]:
+        """Return all indexed L2 chunks for ``doc_id``."""
+        ...
+
+
+@dataclass(frozen=True)
+class ExpansionRequest:
+    """One bounded same-document expansion request."""
+
+    mode: ExpansionMode
+    doc_id: str
+    chunk_index: int | None = None
+    section_path: str | None = None
+    section_paths: tuple[str, ...] = ()
+    reason: str = ""
+
+
+_FIELD_TERMS = (
+    "home resort",
+    "hometown",
+    "born",
+    "discipline",
+    "sponsor",
+    "sponsors",
+    "train",
+    "training",
+    "venue",
+    "located",
+    "hosts",
+    "hosted",
+    "sanctioning",
+    "frequency",
+)
+
+
+def _entity_type_from_doc_id(doc_id: str) -> str:
+    return doc_id.split("/", 1)[0] if "/" in doc_id else ""
+
+
+def _as_tuple(value: Any) -> tuple[str, ...]:
+    if value is None:
+        return ()
+    if isinstance(value, str):
+        return (value,)
+    if isinstance(value, (list, tuple)):
+        return tuple(str(x) for x in value if str(x).strip())
+    return ()
+
+
+def _query_needs_anchor(query: str) -> bool:
+    q = query.lower()
+    return any(t in q for t in _FIELD_TERMS)
+
+
+def plan_expansion_requests(
+    hit: RetrievalHit,
+    query: str,
+    *,
+    modes: tuple[str, ...],
+    window: int,
+    anchor_sections_by_entity_type: Mapping[str, Any],
+) -> list[ExpansionRequest]:
+    """Plan same-document expansion routes for one retrieved hit.
+
+    The planner intentionally returns multiple routes: neighbors work for all
+    chunking strategies, same-section uses Markdown-header metadata when
+    present, and anchors recover entity facts such as athlete ``Summary`` rows.
+    """
+    out: list[ExpansionRequest] = []
+    enabled = {m.strip().lower() for m in modes}
+    if "neighbors" in enabled and hit.chunk_index >= 0 and window > 0:
+        out.append(
+            ExpansionRequest(
+                mode="neighbors",
+                doc_id=hit.doc_id,
+                chunk_index=hit.chunk_index,
+                reason="local_context",
+            )
+        )
+
+    if "same_section" in enabled and hit.section_path:
+        out.append(
+            ExpansionRequest(
+                mode="same_section",
+                doc_id=hit.doc_id,
+                section_path=hit.section_path,
+                reason="same_section",
+            )
+        )
+
+    if "anchor_sections" in enabled and _query_needs_anchor(query):
+        entity_type = _entity_type_from_doc_id(hit.doc_id)
+        anchors = _as_tuple(anchor_sections_by_entity_type.get(entity_type))
+        if anchors:
+            out.append(
+                ExpansionRequest(
+                    mode="anchor_sections",
+                    doc_id=hit.doc_id,
+                    section_paths=anchors,
+                    reason=f"{entity_type}_anchor_sections",
+                )
+            )
+    return out
+
+
+def _vector_to_retrieval(hit: VectorQueryHit) -> RetrievalHit:
+    meta = hit.metadata
+    raw_ci = meta.get("chunk_index", -1)
+    chunk_index = int(raw_ci) if raw_ci is not None else -1
+    return RetrievalHit(
+        chunk_id=hit.id,
+        text=hit.document,
+        doc_id=str(meta.get("doc_id", "")),
+        section_path=str(meta.get("section_path", "")),
+        chunk_index=chunk_index,
+        similarity=chroma_cosine_distance_to_similarity(hit.distance),
+        distance=hit.distance,
+    )
+
+
+def _matches_request(
+    candidate: RetrievalHit,
+    req: ExpansionRequest,
+    *,
+    window: int,
+) -> bool:
+    if candidate.doc_id != req.doc_id:
+        return False
+    if req.mode == "neighbors":
+        if req.chunk_index is None or candidate.chunk_index < 0:
+            return False
+        return 0 < abs(candidate.chunk_index - req.chunk_index) <= window
+    if req.mode == "same_section":
+        return bool(req.section_path) and candidate.section_path == req.section_path
+    if req.mode == "anchor_sections":
+        return candidate.section_path in set(req.section_paths)
+    return False
+
+
+def _penalized_hit(
+    candidate: RetrievalHit,
+    seed: RetrievalHit,
+    penalty: float,
+) -> RetrievalHit:
+    sim = float(seed.similarity) - float(penalty)
+    dist = float(seed.distance) + float(penalty)
+    return RetrievalHit(
+        chunk_id=candidate.chunk_id,
+        text=candidate.text,
+        doc_id=candidate.doc_id,
+        section_path=candidate.section_path,
+        chunk_index=candidate.chunk_index,
+        similarity=sim,
+        distance=dist,
+    )
+
+
+def expand_retrieval_hits(
+    hits: list[RetrievalHit],
+    *,
+    query: str,
+    store: SupportsGetByDocId,
+    config: Mapping[str, Any],
+) -> tuple[list[RetrievalHit], list[RetrievalHit], list[ExpansionRequest]]:
+    """Expand fused retrieval hits with same-document sibling/anchor chunks.
+
+    Returns
+    -------
+    tuple
+        ``(expanded_hits, added_hits, requests)``. ``expanded_hits`` preserves
+        original order and inserts added chunks immediately after their seed.
+    """
+    if not bool(config.get("enabled", False)) or not hits:
+        return hits, [], []
+
+    modes = _as_tuple(config.get("modes", ("neighbors",)))
+    window = max(0, int(config.get("window", 1)))
+    max_per_doc = max(0, int(config.get("max_extra_chunks_per_doc", 3)))
+    max_total = max(0, int(config.get("max_total_extra_chunks", 8)))
+    max_seed_hits = max(1, int(config.get("max_seed_hits", len(hits))))
+    penalty = float(config.get("expansion_score_penalty", 0.05))
+    anchors = config.get("anchor_sections_by_entity_type", {})
+    if not isinstance(anchors, Mapping):
+        anchors = {}
+
+    doc_cache: dict[str, list[RetrievalHit]] = {}
+    requests_by_seed: dict[str, list[ExpansionRequest]] = {}
+    all_requests: list[ExpansionRequest] = []
+    for seed in hits[:max_seed_hits]:
+        reqs = plan_expansion_requests(
+            seed,
+            query,
+            modes=modes,
+            window=window,
+            anchor_sections_by_entity_type=anchors,
+        )
+        requests_by_seed[seed.chunk_id] = reqs
+        all_requests.extend(reqs)
+        if reqs and seed.doc_id not in doc_cache:
+            doc_cache[seed.doc_id] = [
+                _vector_to_retrieval(vh) for vh in store.get_by_doc_id(seed.doc_id)
+            ]
+
+    seen = {h.chunk_id for h in hits}
+    added: list[RetrievalHit] = []
+    expanded: list[RetrievalHit] = []
+    per_doc_counts: dict[str, int] = {}
+
+    for seed in hits:
+        expanded.append(seed)
+        if len(added) >= max_total:
+            continue
+        reqs = requests_by_seed.get(seed.chunk_id, [])
+        if not reqs:
+            continue
+        doc_hits = doc_cache.get(seed.doc_id, [])
+        for req in reqs:
+            if len(added) >= max_total:
+                break
+            for cand in doc_hits:
+                if cand.chunk_id in seen:
+                    continue
+                if per_doc_counts.get(cand.doc_id, 0) >= max_per_doc:
+                    continue
+                if not _matches_request(cand, req, window=window):
+                    continue
+                new_hit = _penalized_hit(cand, seed, penalty)
+                expanded.append(new_hit)
+                added.append(new_hit)
+                seen.add(new_hit.chunk_id)
+                per_doc_counts[new_hit.doc_id] = (
+                    per_doc_counts.get(new_hit.doc_id, 0) + 1
+                )
+                if len(added) >= max_total:
+                    break
+
+    return expanded, added, all_requests
diff --git a/src/vectorstore/chroma_store.py b/src/vectorstore/chroma_store.py
index ec3001c..c7b2fbb 100644
--- a/src/vectorstore/chroma_store.py
+++ b/src/vectorstore/chroma_store.py
@@ -255,3 +255,35 @@ def query(
                 )
             )
         return VectorQueryResult(hits=hits)
+
+    def get_by_doc_id(self, doc_id: str) -> list[VectorQueryHit]:
+        """Return all chunks whose metadata ``doc_id`` matches ``doc_id``.
+
+        Direct Chroma reads do not provide query distances, so returned hits use
+        ``distance=float("inf")``. Callers that inject these rows into a ranked
+        list should assign expansion scores explicitly.
+        """
+        did = str(doc_id).strip()
+        if not did:
+            return []
+        raw = self._collection.get(
+            where={"doc_id": did},
+            include=["documents", "metadatas"],
+        )
+        ids_l = raw.get("ids") or []
+        docs_l = raw.get("documents") or []
+        metas_l = raw.get("metadatas") or []
+        hits: list[VectorQueryHit] = []
+        for i, eid in enumerate(ids_l):
+            doc = docs_l[i] if i < len(docs_l) and docs_l[i] is not None else ""
+            meta = metas_l[i] if i < len(metas_l) and metas_l[i] is not None else {}
+            hits.append(
+                VectorQueryHit(
+                    id=str(eid),
+                    document=str(doc),
+                    distance=float("inf"),
+                    metadata=meta,
+                )
+            )
+        hits.sort(key=lambda h: int(h.metadata.get("chunk_index", 0)))
+        return hits
diff --git a/src/vectorstore/protocol.py b/src/vectorstore/protocol.py
index 3f04b4d..2959440 100644
--- a/src/vectorstore/protocol.py
+++ b/src/vectorstore/protocol.py
@@ -5,7 +5,7 @@
 import numpy as np
 from numpy.typing import NDArray
 
-from .models import VectorQueryResult
+from .models import VectorQueryHit, VectorQueryResult
 
 __all__ = ["VectorStore", "FloatMatrix"]
 
@@ -88,3 +88,19 @@ def query(
             (cosine, L2, inner product).
         """
         ...
+
+    def get_by_doc_id(self, doc_id: str) -> list[VectorQueryHit]:
+        """Return all stored chunks for ``doc_id`` in backend order.
+
+        Parameters
+        ----------
+        doc_id : str
+            Knowledge-base relative source document id.
+
+        Returns
+        -------
+        list of VectorQueryHit
+            Matching rows. Backends may not have similarity distances for direct
+            reads; use ``float("inf")`` or another explicit sentinel.
+        """
+        ...
diff --git a/tests/test_config_merge.py b/tests/test_config_merge.py
index 876f1d0..d3ed936 100644
--- a/tests/test_config_merge.py
+++ b/tests/test_config_merge.py
@@ -59,6 +59,7 @@ def test_merge_app_config_unknown_section_raises() -> None:
         rerank={},
         llm={},
         query_expansion={},
+        document_expansion={},
         generation={},
         logging={},
     )
diff --git a/tests/test_document_expansion.py b/tests/test_document_expansion.py
new file mode 100644
index 0000000..80ad004
--- /dev/null
+++ b/tests/test_document_expansion.py
@@ -0,0 +1,219 @@
+from __future__ import annotations
+
+from dataclasses import replace
+from pathlib import Path
+
+import yaml
+from snow_sports_rag.config import load_config
+from snow_sports_rag.pipeline import RAGPipeline
+from snow_sports_rag.retrieval.document_expansion import (
+    expand_retrieval_hits,
+    plan_expansion_requests,
+)
+from snow_sports_rag.retrieval.models import RetrievalHit
+from snow_sports_rag.vectorstore.models import VectorQueryHit
+
+
+def _hit(
+    doc_id: str,
+    idx: int,
+    *,
+    section: str,
+    text: str = "body",
+    similarity: float = 0.8,
+) -> RetrievalHit:
+    return RetrievalHit(
+        chunk_id=f"{doc_id}::{idx}",
+        text=text,
+        doc_id=doc_id,
+        section_path=section,
+        chunk_index=idx,
+        similarity=similarity,
+        distance=1.0 - similarity,
+    )
+
+
+def _vector_hit(doc_id: str, idx: int, *, section: str, text: str) -> VectorQueryHit:
+    return VectorQueryHit(
+        id=f"{doc_id}::{idx}",
+        document=text,
+        distance=float("inf"),
+        metadata={
+            "doc_id": doc_id,
+            "entity_type": doc_id.split("/", 1)[0],
+            "section_path": section,
+            "chunk_index": idx,
+        },
+    )
+
+
+class _DocStore:
+    def __init__(self, rows: list[VectorQueryHit]) -> None:
+        self.rows = rows
+        self.calls: list[str] = []
+
+    def get_by_doc_id(self, doc_id: str) -> list[VectorQueryHit]:
+        self.calls.append(doc_id)
+        return [r for r in self.rows if r.metadata["doc_id"] == doc_id]
+
+
+def test_route_planner_emits_neighbor_section_and_anchor_requests() -> None:
+    seed = _hit("athletes/Red Gerard.md", 2, section="Career Progression")
+
+    requests = plan_expansion_requests(
+        seed,
+        "What is the home resort for this Olympic medalist?",
+        modes=("neighbors", "same_section", "anchor_sections"),
+        window=1,
+        anchor_sections_by_entity_type={"athletes": ["Summary"]},
+    )
+
+    assert [r.mode for r in requests] == [
+        "neighbors",
+        "same_section",
+        "anchor_sections",
+    ]
+    assert requests[-1].section_paths == ("Summary",)
+
+
+def test_expand_retrieval_hits_adds_anchor_summary_for_field_query() -> None:
+    doc_id = "athletes/Red Gerard.md"
+    seed = _hit(doc_id, 2, section="Career Progression", text="Olympic results")
+    store = _DocStore(
+        [
+            _vector_hit(
+                doc_id,
+                0,
+                section="Summary",
+                text="Home Resort: Copper Mountain. Olympic medals: one gold.",
+            ),
+            _vector_hit(doc_id, 1, section="Career Progression", text="Early career"),
+            _vector_hit(
+                doc_id,
+                2,
+                section="Career Progression",
+                text="Olympic results",
+            ),
+            _vector_hit(doc_id, 3, section="Career Progression", text="Later career"),
+        ]
+    )
+
+    expanded, added, requests = expand_retrieval_hits(
+        [seed],
+        query="Which home resort belongs to the athlete with an Olympic medal?",
+        store=store,
+        config={
+            "enabled": True,
+            "modes": ["neighbors", "same_section", "anchor_sections"],
+            "window": 1,
+            "max_seed_hits": 4,
+            "max_extra_chunks_per_doc": 3,
+            "max_total_extra_chunks": 3,
+            "expansion_score_penalty": 0.05,
+            "anchor_sections_by_entity_type": {"athletes": ["Summary"]},
+        },
+    )
+
+    assert store.calls == [doc_id]
+    assert any(r.mode == "anchor_sections" for r in requests)
+    assert any("Home Resort: Copper Mountain" in h.text for h in added)
+    assert len(added) == 3
+    assert [h.chunk_id for h in expanded][:2] == [seed.chunk_id, f"{doc_id}::1"]
+    assert all(h.similarity == seed.similarity - 0.05 for h in added)
+
+
+def test_expand_retrieval_hits_caps_total_and_per_doc() -> None:
+    doc_id = "athletes/a.md"
+    seed = _hit(doc_id, 2, section="Career")
+    store = _DocStore(
+        [
+            _vector_hit(doc_id, 0, section="Summary", text="home resort"),
+            _vector_hit(doc_id, 1, section="Career", text="left"),
+            _vector_hit(doc_id, 2, section="Career", text="seed"),
+            _vector_hit(doc_id, 3, section="Career", text="right"),
+        ]
+    )
+
+    _, added, _ = expand_retrieval_hits(
+        [seed],
+        query="home resort",
+        store=store,
+        config={
+            "enabled": True,
+            "modes": ["neighbors", "same_section", "anchor_sections"],
+            "window": 1,
+            "max_seed_hits": 1,
+            "max_extra_chunks_per_doc": 1,
+            "max_total_extra_chunks": 10,
+            "anchor_sections_by_entity_type": {"athletes": ["Summary"]},
+        },
+    )
+
+    assert len(added) == 1
+
+
+def test_pipeline_expands_wrong_red_gerard_chunk_before_rerank(tmp_path: Path) -> None:
+    cfg_path = tmp_path / "cfg.yaml"
+    (tmp_path / "kb").mkdir()
+    cfg_path.write_text(
+        yaml.dump(
+            {
+                "knowledge_base_path": "kb",
+                "embedding": {"backend": "fake", "model_name": "fake", "dimension": 8},
+                "generation": {"enabled": False},
+                "query_expansion": {"enabled": False},
+                "rerank": {"enabled": False},
+                "vector_store": {
+                    "backend": "chroma",
+                    "persist_directory": str(tmp_path / "db"),
+                    "collection_name": "snow_sports_kb",
+                    "l1_collection_name": "snow_sports_kb_l1",
+                },
+            }
+        ),
+        encoding="utf-8",
+    )
+    cfg = load_config(cfg_path, base_dir=tmp_path)
+    cfg = replace(
+        cfg,
+        document_expansion={
+            **dict(cfg.document_expansion),
+            "max_extra_chunks_per_doc": 2,
+            "max_total_extra_chunks": 2,
+        },
+    )
+    doc_id = "athletes/Red Gerard.md"
+    seed = _hit(doc_id, 2, section="Career Progression", text="Olympic gold result")
+    pipe = RAGPipeline(cfg)
+    pipe._embedder = object()
+    pipe._l1_store = object()
+    pipe._l2_store = _DocStore(
+        [
+            _vector_hit(
+                doc_id,
+                0,
+                section="Summary",
+                text="Home Resort: Copper Mountain. Olympic medals: one gold.",
+            ),
+            _vector_hit(doc_id, 1, section="Career Progression", text="Early career"),
+            _vector_hit(doc_id, 2, section="Career Progression", text=seed.text),
+        ]
+    )
+    pipe._index_empty = False
+    pipe._retrieve_variant = lambda _query, _preset: ([seed], [doc_id])  # type: ignore[method-assign]
+
+    result = pipe.run(
+        "What is the home resort for the athlete with an Olympic gold medal?",
+        rerank_enabled=False,
+        expansion_enabled=False,
+        generation_enabled=False,
+    )
+
+    assert any(
+        "Home Resort: Copper Mountain" in h.text for h in result.trace.l2_pre_rerank
+    )
+    assert any(
+        "Home Resort: Copper Mountain" in h.text
+        for h in result.trace.document_expansion_added
+    )
+    assert len(result.cards) >= 2
diff --git a/tests/test_vectorstore.py b/tests/test_vectorstore.py
index ba9f514..38499a3 100644
--- a/tests/test_vectorstore.py
+++ b/tests/test_vectorstore.py
@@ -78,6 +78,43 @@ def test_chroma_query_respects_where_doc_id_in(tmp_path: Path) -> None:
     assert out.hits[0].metadata["doc_id"] == "athletes/a.md"
 
 
+def test_chroma_get_by_doc_id_returns_chunks_in_index_order(tmp_path: Path) -> None:
+    store = ChromaVectorStore(tmp_path / "db", "snow_sports_kb")
+    chunks = [
+        Chunk(
+            text="third",
+            doc_id="athletes/a.md",
+            entity_type="athletes",
+            section_path="Career",
+            chunk_index=2,
+        ),
+        Chunk(
+            text="first",
+            doc_id="athletes/a.md",
+            entity_type="athletes",
+            section_path="Summary",
+            chunk_index=0,
+        ),
+        Chunk(
+            text="other",
+            doc_id="athletes/b.md",
+            entity_type="athletes",
+            section_path="Summary",
+            chunk_index=0,
+        ),
+    ]
+    embed = FakeEmbeddingModel(dimension=8, normalize=True)
+    mat = embed.embed_documents([c.text for c in chunks])
+    ids, docs, metas, emb = pack_chunk_upsert(chunks, mat)
+    store.upsert(ids=ids, embeddings=emb, documents=docs, metadatas=metas)
+
+    out = store.get_by_doc_id("athletes/a.md")
+
+    assert [h.metadata["chunk_index"] for h in out] == [0, 2]
+    assert [h.document for h in out] == ["first", "third"]
+    assert store.get_by_doc_id("missing.md") == []
+
+
 def test_chroma_upsert_query_roundtrip(tmp_path: Path) -> None:
     store = ChromaVectorStore(tmp_path / "db", "snow_sports_kb")
     chunks = [

From 87c43fe321e5263d4f366fa9baeb98917511a82c Mon Sep 17 00:00:00 2001
From: Morgan-Sell <morganpsell@gmail.com>
Date: Fri, 24 Apr 2026 18:50:19 -0700
Subject: [PATCH 2/8] fixed formatting error

---
 tests/test_gold_qa.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/tests/test_gold_qa.py b/tests/test_gold_qa.py
index 417bf41..ff1f18b 100644
--- a/tests/test_gold_qa.py
+++ b/tests/test_gold_qa.py
@@ -6,8 +6,11 @@
 from pathlib import Path
 
 import pytest
-
-from snow_sports_rag.evaluation import GoldItem, default_gold_qa_path, load_gold_qa
+from snow_sports_rag.evaluation import (
+    GoldItem,
+    default_gold_qa_path,
+    load_gold_qa
+)
 
 
 def test_default_gold_file_exists() -> None:

From e86191c3e9f8464a1907ea2d294d747822f58de0 Mon Sep 17 00:00:00 2001
From: Morgan-Sell <morganpsell@gmail.com>
Date: Fri, 24 Apr 2026 18:54:10 -0700
Subject: [PATCH 3/8] fixed formatting error - take 2

---
 tests/test_gold_qa.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/test_gold_qa.py b/tests/test_gold_qa.py
index ff1f18b..106ba0a 100644
--- a/tests/test_gold_qa.py
+++ b/tests/test_gold_qa.py
@@ -7,9 +7,9 @@
 
 import pytest
 from snow_sports_rag.evaluation import (
-    GoldItem,
     default_gold_qa_path,
-    load_gold_qa
+    load_gold_qa,
+    GoldItem,
 )
 
 

From dd5ed5a9f4734ddcee4671969a3b2df26c930eef Mon Sep 17 00:00:00 2001
From: Morgan-Sell <morganpsell@gmail.com>
Date: Fri, 24 Apr 2026 18:56:48 -0700
Subject: [PATCH 4/8] fixed formatting error - take 3

---
 tests/test_gold_qa.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_gold_qa.py b/tests/test_gold_qa.py
index 106ba0a..a97ea07 100644
--- a/tests/test_gold_qa.py
+++ b/tests/test_gold_qa.py
@@ -8,8 +8,8 @@
 import pytest
 from snow_sports_rag.evaluation import (
     default_gold_qa_path,
-    load_gold_qa,
     GoldItem,
+    load_gold_qa,
 )
 
 

From 7e4fbba8ac91d5f253e7108d60a3950f4e52668a Mon Sep 17 00:00:00 2001
From: Morgan-Sell <morganpsell@gmail.com>
Date: Fri, 24 Apr 2026 19:03:52 -0700
Subject: [PATCH 5/8] fixed formatting error - take 4

---
 tests/test_gold_qa.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tests/test_gold_qa.py b/tests/test_gold_qa.py
index a97ea07..7329413 100644
--- a/tests/test_gold_qa.py
+++ b/tests/test_gold_qa.py
@@ -6,9 +6,10 @@
 from pathlib import Path
 
 import pytest
+
 from snow_sports_rag.evaluation import (
-    default_gold_qa_path,
     GoldItem,
+    default_gold_qa_path,
     load_gold_qa,
 )
 

From b5f8f967d7bdbc6806c0b38155e81b707e38a974 Mon Sep 17 00:00:00 2001
From: Morgan-Sell <morganpsell@gmail.com>
Date: Fri, 24 Apr 2026 19:10:11 -0700
Subject: [PATCH 6/8] changed ruff version on pyproject.toml to match uv.lock

---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 6298577..347dd25 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -45,7 +45,7 @@ packages = [
 [dependency-groups]
 dev = [
     "pytest>=8.0",
-    "ruff>=0.8",
+    "ruff==0.15.10",
 ]
 
 [tool.pytest.ini_options]

From 693bf86cc0d6c25a674f81c764adea0157aa659e Mon Sep 17 00:00:00 2001
From: Morgan-Sell <morganpsell@gmail.com>
Date: Fri, 24 Apr 2026 19:18:11 -0700
Subject: [PATCH 7/8] added ruff debug logic to ci.yaml

---
 .github/workflows/ci.yml | 7 ++++++-
 uv.lock                  | 3 ++-
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 5d60433..ebf96d0 100644
--- a/.github/workflows/ci.yml
+++ b/.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
diff --git a/uv.lock b/uv.lock
index cccc3e6..25477ac 100644
--- a/uv.lock
+++ b/uv.lock
@@ -668,6 +668,7 @@ dependencies = [
     { name = "gradio-client" },
     { name = "typer" },
 ]
+sdist = { url = "https://files.pythonhosted.org/packages/ce/86/c9694b7cfada5780e75769e60dc161a161f4dd7fc91b61db5e3a3338bef9/hf_gradio-0.4.1.tar.gz", hash = "sha256:a017d942618f0d495a58ee4563047fa04bef614c00e0cb789a9a6d0633cffa7b", size = 6560, upload-time = "2026-04-22T14:01:32.334Z" }
 wheels = [
     { url = "https://files.pythonhosted.org/packages/30/2d/afff2ee87e75d8eb85c92bb8cf0e15b05c23c2ebd8fd8dec781d8601ed7f/hf_gradio-0.4.1-py3-none-any.whl", hash = "sha256:76b8cb8be6abe62d74c1ad2d35b42f0629db89aa9e1a8d033cecfe7c856eeab3", size = 4482, upload-time = "2026-04-17T19:53:31.827Z" },
 ]
@@ -2679,7 +2680,7 @@ requires-dist = [
 [package.metadata.requires-dev]
 dev = [
     { name = "pytest", specifier = ">=8.0" },
-    { name = "ruff", specifier = ">=0.8" },
+    { name = "ruff", specifier = "==0.15.10" },
 ]
 
 [[package]]

From c686091d30f3174ad9f8fcd501427ca23b3b275e Mon Sep 17 00:00:00 2001
From: Morgan-Sell <morganpsell@gmail.com>
Date: Fri, 24 Apr 2026 19:30:18 -0700
Subject: [PATCH 8/8] snow_sports_rag is now a known-first party, so added
 blank line between imports

---
 pyproject.toml                   | 3 +++
 tests/test_chunking.py           | 1 +
 tests/test_config.py             | 1 +
 tests/test_config_merge.py       | 1 +
 tests/test_document_expansion.py | 1 +
 tests/test_embedding.py          | 1 +
 tests/test_fusion.py             | 1 +
 tests/test_generation.py         | 1 +
 tests/test_llm.py                | 1 +
 tests/test_loader.py             | 1 +
 tests/test_pipeline.py           | 1 +
 tests/test_query_expansion.py    | 1 +
 tests/test_rerank.py             | 1 +
 tests/test_retrieval_baseline.py | 1 +
 tests/test_sweep_runner.py       | 1 +
 tests/test_vectorstore.py        | 1 +
 16 files changed, 18 insertions(+)

diff --git a/pyproject.toml b/pyproject.toml
index 347dd25..bf05152 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -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"]
\ No newline at end of file
diff --git a/tests/test_chunking.py b/tests/test_chunking.py
index 4958be6..15f8275 100644
--- a/tests/test_chunking.py
+++ b/tests/test_chunking.py
@@ -3,6 +3,7 @@
 from pathlib import Path
 
 import pytest
+
 from snow_sports_rag.chunking import (
     FixedWindowChunkStrategy,
     MarkdownHeaderChunkStrategy,
diff --git a/tests/test_config.py b/tests/test_config.py
index fb84046..4ae79e8 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -4,6 +4,7 @@
 
 import pytest
 import yaml
+
 from snow_sports_rag.config import load_config
 
 
diff --git a/tests/test_config_merge.py b/tests/test_config_merge.py
index d3ed936..3e27b84 100644
--- a/tests/test_config_merge.py
+++ b/tests/test_config_merge.py
@@ -5,6 +5,7 @@
 from pathlib import Path
 
 import pytest
+
 from snow_sports_rag.config.loader import AppConfig, load_config
 from snow_sports_rag.evaluation.config_merge import (
     flat_overrides_to_nested,
diff --git a/tests/test_document_expansion.py b/tests/test_document_expansion.py
index 80ad004..9bbd1b2 100644
--- a/tests/test_document_expansion.py
+++ b/tests/test_document_expansion.py
@@ -4,6 +4,7 @@
 from pathlib import Path
 
 import yaml
+
 from snow_sports_rag.config import load_config
 from snow_sports_rag.pipeline import RAGPipeline
 from snow_sports_rag.retrieval.document_expansion import (
diff --git a/tests/test_embedding.py b/tests/test_embedding.py
index a782c4e..4020491 100644
--- a/tests/test_embedding.py
+++ b/tests/test_embedding.py
@@ -2,6 +2,7 @@
 
 import numpy as np
 import pytest
+
 from snow_sports_rag.embedding import (
     FakeEmbeddingModel,
     SentenceTransformerEmbeddingModel,
diff --git a/tests/test_fusion.py b/tests/test_fusion.py
index a58c775..2af3002 100644
--- a/tests/test_fusion.py
+++ b/tests/test_fusion.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import pytest
+
 from snow_sports_rag.retrieval import (
     fuse_retrieval_hits_max_score,
     fuse_retrieval_hits_rrf,
diff --git a/tests/test_generation.py b/tests/test_generation.py
index f289a1a..f878db9 100644
--- a/tests/test_generation.py
+++ b/tests/test_generation.py
@@ -4,6 +4,7 @@
 from unittest.mock import patch
 
 import pytest
+
 from snow_sports_rag.generation import (
     DEFAULT_REFUSAL_MESSAGE,
     DEFAULT_SYSTEM_PROMPT,
diff --git a/tests/test_llm.py b/tests/test_llm.py
index 44b1112..be86770 100644
--- a/tests/test_llm.py
+++ b/tests/test_llm.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import pytest
+
 from snow_sports_rag.llm import FakeLLMClient, llm_client_from_config
 from snow_sports_rag.llm.openai_compatible import (
     OpenAICompatibleLLMClient,
diff --git a/tests/test_loader.py b/tests/test_loader.py
index 957b4a0..d271329 100644
--- a/tests/test_loader.py
+++ b/tests/test_loader.py
@@ -3,6 +3,7 @@
 from pathlib import Path
 
 import yaml
+
 from snow_sports_rag.config import load_config
 from snow_sports_rag.ingest import KnowledgeBaseLoader
 
diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py
index 3e0b80d..44f4bfe 100644
--- a/tests/test_pipeline.py
+++ b/tests/test_pipeline.py
@@ -5,6 +5,7 @@
 
 import pytest
 import yaml
+
 from snow_sports_rag.chunking import MarkdownHeaderChunkStrategy
 from snow_sports_rag.config import load_config
 from snow_sports_rag.embedding import FakeEmbeddingModel
diff --git a/tests/test_query_expansion.py b/tests/test_query_expansion.py
index 7c4ef40..3ce8d08 100644
--- a/tests/test_query_expansion.py
+++ b/tests/test_query_expansion.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import pytest
+
 from snow_sports_rag.llm import FakeLLMClient
 from snow_sports_rag.retrieval import QueryExpander
 from snow_sports_rag.retrieval.models import RetrievalHit
diff --git a/tests/test_rerank.py b/tests/test_rerank.py
index 165e635..cbd6974 100644
--- a/tests/test_rerank.py
+++ b/tests/test_rerank.py
@@ -5,6 +5,7 @@
 
 import numpy as np
 import pytest
+
 from snow_sports_rag.rerank import (
     CrossEncoderReranker,
     IdentityReranker,
diff --git a/tests/test_retrieval_baseline.py b/tests/test_retrieval_baseline.py
index 7e619e3..20153e0 100644
--- a/tests/test_retrieval_baseline.py
+++ b/tests/test_retrieval_baseline.py
@@ -4,6 +4,7 @@
 
 import pytest
 import yaml
+
 from snow_sports_rag.chunking import (
     MarkdownHeaderChunkStrategy,
     chunk_strategy_from_config,
diff --git a/tests/test_sweep_runner.py b/tests/test_sweep_runner.py
index f6f203d..78842d4 100644
--- a/tests/test_sweep_runner.py
+++ b/tests/test_sweep_runner.py
@@ -6,6 +6,7 @@
 from pathlib import Path
 
 import yaml
+
 from snow_sports_rag.config.loader import load_config
 from snow_sports_rag.evaluation.gold import load_gold_qa
 from snow_sports_rag.evaluation.sweep import SweepConfig, run_sweep_retrieval
diff --git a/tests/test_vectorstore.py b/tests/test_vectorstore.py
index 38499a3..b231dbb 100644
--- a/tests/test_vectorstore.py
+++ b/tests/test_vectorstore.py
@@ -4,6 +4,7 @@
 
 import numpy as np
 import pytest
+
 from snow_sports_rag.chunking.models import Chunk
 from snow_sports_rag.config import load_config
 from snow_sports_rag.embedding import FakeEmbeddingModel