fix context expansion: respect section boundaries, remove max_context_items

docs/configuration/index.md

@@ -99,7 +99,6 @@ research:
 
 search:
   limit: 10                    # Default number of results to return
-  max_context_items: 10        # Maximum items in expanded context
   max_context_chars: 10000     # Maximum characters in expanded context
   vector_index_metric: cosine  # cosine, l2, or dot
   vector_refine_factor: 30

docs/configuration/qa-research.md

@@ -7,12 +7,10 @@ Configure search behavior and context expansion:
 ```yaml
 search:
   limit: 10                    # Default number of results to return
-  max_context_items: 10        # Maximum items in expanded context
   max_context_chars: 10000     # Maximum characters in expanded context
 ```
 
 - **limit**: Default number of search results to return when no limit is specified. Used by CLI, MCP server, QA, and research workflows. Default: 10
-- **max_context_items**: Limits how many document items (paragraphs, list items, etc.) can be included in expanded context. Default: 10.
 - **max_context_chars**: Hard limit on total characters in expanded content. Default: 10000.
 
 Context expansion is automatic and section-aware. For structured documents (with section headers), expansion includes the entire section containing the match. For sections that exceed the budget or are too small (e.g., a title+authors area), expansion grows outward item-by-item from the match center, skipping noise labels (footnotes, page headers) — this naturally crosses into adjacent sections until the budget is filled. For unstructured documents, expansion grows outward item-by-item. Results without `doc_item_refs` (e.g., custom chunks passed to `import_document`) pass through unexpanded.

docs/configuration/storage.md

@@ -112,7 +112,7 @@ search:
   vector_refine_factor: 30     # Re-ranking factor for accuracy
 ```
 
-For search behavior settings (`limit`, `max_context_items`, `max_context_chars`), see [QA and Research](qa-research.md#search-settings).
+For search behavior settings (`limit`, `max_context_chars`), see [QA and Research](qa-research.md#search-settings).
 
 - **vector_index_metric**: Distance metric for vector similarity:
   - `cosine`: Cosine similarity (default, best for most embeddings)

docs/python.md

@@ -376,7 +376,6 @@ Context expansion is automatic and section-aware. For structured documents (with
 
 Configuration:
 
-- **search.max_context_items**: Maximum items in expanded context. Default: 10.
 - **search.max_context_chars**: Maximum characters in expanded context. Default: 10000.
 
 **Smart Merging**: When expanded results overlap within the same document, they are automatically merged into a single result with continuous content and the highest relevance score.

docs/tuning.md

@@ -26,7 +26,7 @@ When configured, a cross-encoder reranker re-scores 10x the requested candidates
 
 `limit` controls how many results reach the LLM. More candidates improve recall but increase token usage. See [Search Settings](configuration/qa-research.md#search-settings).
 
-Context expansion is automatic and section-aware — search results are expanded to include surrounding content from the same document section. For structured documents, expansion stays within section boundaries and filters noise (footnotes, page headers). For unstructured documents, expansion grows outward until the character budget is filled. `max_context_items` and `max_context_chars` cap expansion to prevent context bloat.
+Context expansion is automatic and section-aware — search results are expanded to include surrounding content from the same document section. For structured documents, expansion stays within section boundaries and filters noise (footnotes, page headers). For unstructured documents, expansion grows outward until the character budget is filled. `max_context_chars` caps expansion to prevent context bloat.
 
 ## Tuning Generation
 

evaluations/evaluations/benchmark.py

@@ -49,7 +49,6 @@ def build_experiment_metadata(
         "embedder_dim": config.embeddings.model.vector_dim,
         "chunk_size": config.processing.chunk_size,
         "search_limit": config.search.limit,
-        "max_context_items": config.search.max_context_items,
         "max_context_chars": config.search.max_context_chars,
         "rerank_provider": config.reranking.model.provider
         if config.reranking.model

haiku_rag_slim/haiku/rag/client.py

@@ -1105,7 +1105,6 @@ class HaikuRAG:
         """
         from haiku.rag.context import expand_with_items
 
-        max_items = self._config.search.max_context_items
         max_chars = self._config.search.max_context_chars
 
         # Group by document_id for efficient processing
@@ -1132,7 +1131,6 @@ class HaikuRAG:
                 self.document_item_repository,
                 doc_id,
                 doc_results,
-                max_items,
                 max_chars,
             )
             expanded_results.extend(expanded)
@@ -1272,9 +1270,9 @@ class HaikuRAG:
     async def visualize_chunk(self, chunk: Chunk) -> list:
         """Render page images with bounding box highlights for a chunk.
 
-        Gets the DoclingDocument from the chunk's document, resolves bounding boxes
+        Expands the chunk's context to find the full section, then resolves
-        from chunk metadata, and renders all pages that contain bounding boxes with
+        bounding boxes from all items in the expanded range. This ensures
-        yellow/orange highlight overlays.
+        visualization covers all pages the expanded content spans.
 
         Args:
             chunk: The chunk to visualize.
@@ -1287,6 +1285,8 @@ class HaikuRAG:
 
         from PIL import ImageDraw
 
+        from haiku.rag.store.models.chunk import ChunkMetadata
+
         # Get the document structure (from cache if available)
         if not chunk.document_id:
             return []
@@ -1299,9 +1299,23 @@ class HaikuRAG:
         if not docling_doc:
             return []
 
-        # Resolve bounding boxes from chunk metadata
+        # Expand context to get all doc_item_refs in the section
         chunk_meta = chunk.get_chunk_metadata()
-        bounding_boxes = chunk_meta.resolve_bounding_boxes(docling_doc)
+        if chunk_meta.doc_item_refs:
+            search_result = SearchResult(
+                content=chunk.content,
+                score=1.0,
+                chunk_id=chunk.id,
+                document_id=chunk.document_id,
+                doc_item_refs=chunk_meta.doc_item_refs,
+                page_numbers=chunk_meta.page_numbers,
+            )
+            expanded = await self.expand_context([search_result])
+            refs = expanded[0].doc_item_refs if expanded else chunk_meta.doc_item_refs
+            meta = ChunkMetadata(doc_item_refs=refs)
+        else:
+            meta = chunk_meta
+        bounding_boxes = meta.resolve_bounding_boxes(docling_doc)
         if not bounding_boxes:
             return []
 

haiku_rag_slim/haiku/rag/config/models.py

@@ -174,7 +174,6 @@ class ProcessingConfig(BaseModel):
 
 class SearchConfig(BaseModel):
     limit: int = 10
-    max_context_items: int = 10
     max_context_chars: int = 10000
     vector_index_metric: Literal["cosine", "l2", "dot"] = "cosine"
     vector_refine_factor: int = 30

haiku_rag_slim/haiku/rag/context.py

@@ -6,12 +6,16 @@ the document_items table. The algorithm adapts to document structure:
 For STRUCTURED documents (containing section_header or title labels):
   1. Resolve matched doc_item_refs to positions in the items table
   2. Find section boundaries around each match (section_header/title labels)
-  3. If the section fits within the budget, include it entirely
+  3. If the section fits within the char budget, include it entirely
-  4. If the section exceeds the budget, OR the section is too small (under
+  4. If the section exceeds the char budget, expand item-by-item from the
-     20% of max_context_chars), expand item-by-item from the match center
+     match center outward, bounded by section edges
-     outward, skipping noise labels. This lets small sections (e.g., a
+  5. If the section is too small (under 20% of max_context_chars), expand
-     title+authors area) grow into the next section's content.
+     item-by-item crossing into adjacent sections until the budget is filled.
-  5. Merge overlapping ranges from multiple results in the same document
+     This lets small sections (e.g., title+authors) grow into neighboring
+     content.
+  6. Merge overlapping ranges from multiple results in the same document.
+     Adjacent but non-overlapping ranges stay separate to preserve section
+     independence.
 
 For UNSTRUCTURED documents (no section headers):
   Expand outward item-by-item from the match center until the character
@@ -20,7 +24,6 @@ For UNSTRUCTURED documents (no section headers):
 
 In both cases:
   - max_context_chars caps total characters per expanded result
-  - max_context_items caps total items per expanded result
   - Noise labels (footnote, page_header, page_footer, document_index) are
     excluded from content AND budget counting in structured documents
   - Results without doc_item_refs pass through unexpanded
@@ -42,7 +45,7 @@ _MIN_SECTION_BUDGET_RATIO = 0.2
 def _merge_ranges(
     ranges: list[tuple[int, int, SearchResult]],
 ) -> list[tuple[int, int, list[SearchResult]]]:
-    """Merge overlapping or adjacent ranges."""
+    """Merge overlapping ranges. Adjacent but non-overlapping ranges stay separate."""
     if not ranges:
         return []
 
@@ -55,7 +58,7 @@ def _merge_ranges(
     )
 
     for min_idx, max_idx, result in sorted_ranges[1:]:
-        if cur_max >= min_idx - 1:  # Overlapping or adjacent
+        if cur_max >= min_idx:  # Truly overlapping
             cur_max = max(cur_max, max_idx)
             cur_results.append(result)
         else:
@@ -69,27 +72,32 @@ def _merge_ranges(
 def _expand_outward(
     items: list[DocumentItem],
     center_idx: int,
-    max_items: int,
     max_chars: int,
     skip_noise: bool = False,
+    lo_bound: int = 0,
+    hi_bound: int | None = None,
 ) -> tuple[int, int]:
-    """Expand item-by-item outward from center until budget is filled.
+    """Expand item-by-item outward from center until char budget is filled.
 
     When skip_noise is True, noise labels are excluded from char counting
     (used in structured documents so footnotes don't consume budget).
+
+    lo_bound and hi_bound constrain expansion (e.g., to section edges).
     """
+    if hi_bound is None:
+        hi_bound = len(items) - 1
     lo = hi = center_idx
     center_is_noise = skip_noise and items[center_idx].label in _NOISE_LABELS
     char_count = 0 if center_is_noise else len(items[center_idx].text)
 
-    while char_count < max_chars and hi - lo + 1 < max_items:
+    while char_count < max_chars:
         grew = False
-        if lo > 0:
+        if lo > lo_bound:
             lo -= 1
             if not (skip_noise and items[lo].label in _NOISE_LABELS):
                 char_count += len(items[lo].text)
             grew = True
-        if hi < len(items) - 1 and char_count < max_chars:
+        if hi < hi_bound and char_count < max_chars:
             hi += 1
             if not (skip_noise and items[hi].label in _NOISE_LABELS):
                 char_count += len(items[hi].text)
@@ -104,7 +112,6 @@ def _find_expansion_range(
     items: list[DocumentItem],
     matched_positions: set[int],
     has_sections: bool,
-    max_items: int,
     max_chars: int,
 ) -> tuple[int, int]:
     """Find the expansion range for matched positions within a window of items."""
@@ -113,7 +120,7 @@ def _find_expansion_range(
     center_idx = matched_indices[len(matched_indices) // 2]
 
     if not has_sections:
-        return _expand_outward(items, center_idx, max_items, max_chars)
+        return _expand_outward(items, center_idx, max_chars)
 
     # Build section spans: [(start_idx, end_idx), ...]
     headers = [
@@ -140,23 +147,34 @@ def _find_expansion_range(
         if items[i].label not in _NOISE_LABELS
     )
 
-    # Section fits nicely in the budget — return it as-is
     min_useful = int(max_chars * _MIN_SECTION_BUDGET_RATIO)
-    if min_useful <= sec_chars <= max_chars and sec_end - sec_start + 1 <= max_items:
+
+    if sec_chars <= max_chars and sec_chars >= min_useful:
+        # Section fits in char budget — return it regardless of item count
         return (items[sec_start].position, items[sec_end].position)
 
-    # Section is too large or too small — expand item-by-item from center.
+    if sec_chars > max_chars:
-    # For too-large sections this stays within budget.
+        # Section too large — expand outward bounded by section edges
-    # For too-small sections (e.g., title+authors) this naturally grows
+        return _expand_outward(
-    # into adjacent sections until the budget is filled.
+            items,
-    return _expand_outward(items, center_idx, max_items, max_chars, skip_noise=True)
+            center_idx,
+            max_chars,
+            skip_noise=True,
+            lo_bound=sec_start,
+            hi_bound=sec_end,
+        )
+
+    # Section too small (e.g., title+authors) — expand across boundaries
+    return _expand_outward(items, center_idx, max_chars, skip_noise=True)
+
+
+_WINDOW_MARGIN = 100
 
 
 async def expand_with_items(
     document_item_repository: DocumentItemRepository,
     document_id: str,
     results: list[SearchResult],
-    max_items: int,
     max_chars: int,
 ) -> list[SearchResult]:
     """Expand results using the document_items table."""
@@ -172,7 +190,7 @@ async def expand_with_items(
     # wide enough to find section boundaries (the nearest section_header/title
     # above and below the match).
     all_positions = sorted(ref_positions.values())
-    window_margin = max_items * 10
+    window_margin = _WINDOW_MARGIN
     window_start = max(0, min(all_positions) - window_margin)
     window_end = max(all_positions) + window_margin
     window_items = await document_item_repository.get_items_in_range(
@@ -194,9 +212,7 @@ async def expand_with_items(
             passthrough.append(result)
             continue
 
-        lo, hi = _find_expansion_range(
+        lo, hi = _find_expansion_range(window_items, matched, has_sections, max_chars)
-            window_items, matched, has_sections, max_items, max_chars
-        )
         ranges.append((lo, hi, result))
 
     merged = _merge_ranges(ranges)

tests/test_context.py

@@ -49,11 +49,12 @@ class TestMergeRanges:
         assert len(merged) == 1
         assert merged[0] == (0, 15, [r1, r2])
 
-    def test_adjacent(self):
+    def test_adjacent_stay_separate(self):
         r1, r2 = _result(), _result()
         merged = _merge_ranges([(0, 5, r1), (6, 10, r2)])
-        assert len(merged) == 1
+        assert len(merged) == 2
-        assert merged[0] == (0, 10, [r1, r2])
+        assert merged[0] == (0, 5, [r1])
+        assert merged[1] == (6, 10, [r2])
 
     def test_sorts_by_position(self):
         r1, r2 = _result(), _result()
@@ -65,7 +66,7 @@ class TestMergeRanges:
 class TestExpandOutward:
     def test_basic_expansion(self):
         items = [_item(i, text=f"{'x' * 100}") for i in range(10)]
-        lo, hi = _expand_outward(items, 5, max_items=10, max_chars=500)
+        lo, hi = _expand_outward(items, 5, max_chars=500)
         assert lo <= 5
         assert hi >= 5
         total = sum(
@@ -76,16 +77,9 @@ class TestExpandOutward:
         # Should be around 500 chars (may overshoot by one item)
         assert total >= 400
 
-    def test_respects_max_items(self):
-        items = [_item(i, text="x") for i in range(100)]
-        lo, hi = _expand_outward(items, 50, max_items=5, max_chars=999999)
-        count = hi - lo + 1
-        # May overshoot by 1-2 items due to alternating expansion
-        assert count <= 7
-
     def test_respects_max_chars(self):
         items = [_item(i, text=f"{'x' * 200}") for i in range(20)]
-        lo, hi = _expand_outward(items, 10, max_items=999, max_chars=500)
+        lo, hi = _expand_outward(items, 10, max_chars=500)
         total = sum(
             len(items[i].text)
             for i in range(lo, hi + 1)
@@ -96,12 +90,12 @@ class TestExpandOutward:
 
     def test_center_at_start(self):
         items = [_item(i) for i in range(10)]
-        lo, hi = _expand_outward(items, 0, max_items=5, max_chars=999999)
+        lo, hi = _expand_outward(items, 0, max_chars=999999)
         assert lo == 0
 
     def test_center_at_end(self):
         items = [_item(i) for i in range(10)]
-        lo, hi = _expand_outward(items, 9, max_items=5, max_chars=999999)
+        lo, hi = _expand_outward(items, 9, max_chars=999999)
         assert hi == 9
 
     def test_skip_noise_excludes_from_char_count(self):
@@ -113,7 +107,7 @@ class TestExpandOutward:
             _item(4, label="footnote", text="f" * 5000),
             _item(5, text="d" * 100),
         ]
-        lo, hi = _expand_outward(items, 2, max_items=10, max_chars=500, skip_noise=True)
+        lo, hi = _expand_outward(items, 2, max_chars=500, skip_noise=True)
         # Footnotes (5000 chars each) should NOT count toward budget
         # So we should expand past them
         assert lo <= 0
@@ -125,11 +119,17 @@ class TestExpandOutward:
             _item(1, label="document_index", text="x" * 10000),
             _item(2, text="b" * 200),
         ]
-        lo, hi = _expand_outward(items, 1, max_items=10, max_chars=500, skip_noise=True)
+        lo, hi = _expand_outward(items, 1, max_chars=500, skip_noise=True)
         # Center is noise, should start at 0 chars and expand outward
         assert lo == 0
         assert hi == 2
 
+    def test_respects_bounds(self):
+        items = [_item(i, text="x" * 100) for i in range(20)]
+        lo, hi = _expand_outward(items, 10, max_chars=999999, lo_bound=8, hi_bound=12)
+        assert lo == 8
+        assert hi == 12
+
 
 class TestFindExpansionRange:
     def _structured_items(self):
@@ -146,36 +146,47 @@ class TestFindExpansionRange:
 
     def test_structured_returns_section(self):
         items = self._structured_items()
-        lo, hi = _find_expansion_range(
+        lo, hi = _find_expansion_range(items, {1}, has_sections=True, max_chars=5000)
-            items, {1}, has_sections=True, max_items=20, max_chars=5000
-        )
         # Should return the Introduction section (items 0-3)
         assert lo == 0
         assert hi == 3
 
     def test_structured_different_section(self):
         items = self._structured_items()
-        lo, hi = _find_expansion_range(
+        lo, hi = _find_expansion_range(items, {5}, has_sections=True, max_chars=5000)
-            items, {5}, has_sections=True, max_items=20, max_chars=5000
-        )
         # Should return the Methods section (items 4-6)
         assert lo == 4
         assert hi == 6
 
-    def test_structured_large_section_falls_back_to_outward(self):
+    def test_structured_large_section_bounded_by_section(self):
         items = [
             _item(0, label="section_header", text="Big Section"),
         ] + [_item(i, text="x" * 1000) for i in range(1, 20)]
         # Section has 19 * 1000 = 19000 chars, way over 5000 budget
-        lo, hi = _find_expansion_range(
+        lo, hi = _find_expansion_range(items, {10}, has_sections=True, max_chars=5000)
-            items, {10}, has_sections=True, max_items=50, max_chars=5000
+        # Should NOT return the full section, but should stay within it
-        )
-        # Should NOT return the full section
         total = sum(
             len(items[i].text) for i in range(lo, hi + 1) if items[i].position >= lo
         )
         assert total < 10000
 
+    def test_structured_section_with_many_items_returned_whole(self):
+        """A section that fits in char budget is returned even with many items."""
+        items = (
+            [
+                _item(0, label="section_header", text="Section"),
+            ]
+            + [_item(i, text="x" * 200) for i in range(1, 20)]
+            + [
+                _item(20, label="section_header", text="Next"),
+            ]
+        )
+        # Section has 19 * 200 = 3800 chars + header, under 5000 and over min_useful
+        lo, hi = _find_expansion_range(items, {10}, has_sections=True, max_chars=5000)
+        # Should return entire section despite 20 items
+        assert lo == 0
+        assert hi == 19
+
     def test_structured_small_section_expands_outward(self):
         items = [
             _item(0, label="title", text="Paper Title"),
@@ -186,26 +197,20 @@ class TestFindExpansionRange:
             _item(5, text="Intro content. " * 50),
         ]
         # Title section (items 0-1) is tiny (~25 chars) < 20% of 5000
-        lo, hi = _find_expansion_range(
+        lo, hi = _find_expansion_range(items, {0}, has_sections=True, max_chars=5000)
-            items, {0}, has_sections=True, max_items=20, max_chars=5000
-        )
         # Should expand past the title section into the abstract
         assert hi >= 3
 
     def test_unstructured_expands_outward(self):
         items = [_item(i, text=f"Paragraph {i}. " * 10) for i in range(10)]
-        lo, hi = _find_expansion_range(
+        lo, hi = _find_expansion_range(items, {5}, has_sections=False, max_chars=5000)
-            items, {5}, has_sections=False, max_items=20, max_chars=5000
-        )
         assert lo < 5
         assert hi > 5
 
     def test_multiple_matched_positions_uses_center(self):
         items = [_item(i, text="x" * 100) for i in range(20)]
-        # Match at positions 3 and 7, center should be index for position 5 (median)
+        # Use a char budget that forces partial expansion so center matters
-        lo, hi = _find_expansion_range(
+        lo, hi = _find_expansion_range(items, {3, 7}, has_sections=False, max_chars=500)
-            items, {3, 7}, has_sections=False, max_items=5, max_chars=999999
-        )
         center = (lo + hi) // 2
         # Center should be around position 5
         assert 3 <= center <= 7
@@ -219,9 +224,7 @@ class TestFindExpansionRange:
         ]
         # Section non-noise chars: ~260 chars (items 0,1,3). Under 5000 budget.
         # The footnote's 10000 chars should NOT count.
-        lo, hi = _find_expansion_range(
+        lo, hi = _find_expansion_range(items, {1}, has_sections=True, max_chars=5000)
-            items, {1}, has_sections=True, max_items=20, max_chars=5000
-        )
         # Should return full section (it fits in budget excluding noise)
         assert lo == 0
         assert hi == 3
@@ -233,9 +236,7 @@ class TestFindExpansionRange:
             _item(2, label="section_header", text="First Section"),
             _item(3, text="Section content."),
         ]
-        lo, hi = _find_expansion_range(
+        lo, hi = _find_expansion_range(items, {0}, has_sections=True, max_chars=5000)
-            items, {0}, has_sections=True, max_items=20, max_chars=5000
-        )
         # Match is in preamble section (items 0-1), which is small
         # Should expand outward into the first section
         assert hi >= 2
@@ -262,7 +263,7 @@ class TestExpandWithItems:
                 doc_item_refs=["#/texts/999999"],
             )
             expanded = await expand_with_items(
-                rag.document_item_repository, doc.id, [result], 10, 5000
+                rag.document_item_repository, doc.id, [result], 5000
             )
             assert len(expanded) == 1
             assert expanded[0].content == "original"
@@ -312,7 +313,7 @@ class TestExpandWithItems:
                 doc_item_refs=["#/texts/1"],
             )
             expanded = await expand_with_items(
-                rag.document_item_repository, "doc-1", [result], 10, 5000
+                rag.document_item_repository, "doc-1", [result], 5000
             )
             assert len(expanded) == 1
             # The TOC section's only non-header item is document_index (noise).
@@ -376,7 +377,7 @@ class TestExpandWithItems:
                 doc_item_refs=["#/texts/1", "#/texts/2", "#/texts/3", "#/texts/4"],
             )
             expanded = await expand_with_items(
-                rag.document_item_repository, "doc-1", [result], 10, 5000
+                rag.document_item_repository, "doc-1", [result], 5000
             )
             assert len(expanded) == 1
             # Expansion produces "Steps\n\nClick\n\n+\n\nAdd a New Service" = 38 chars

tests/test_context_enhancement.py

@@ -80,7 +80,6 @@ def small_chunk_config() -> AppConfig:
     """Config with small chunk size to force splitting."""
     config = AppConfig()
     config.processing.chunk_size = 32
-    config.search.max_context_items = 25
     config.search.max_context_chars = 10000
     return config
 
@@ -305,33 +304,6 @@ async def test_format_for_agent_output(temp_db_path, small_chunk_config):
         assert "Content:" in formatted
 
 
-@pytest.mark.vcr()
-async def test_max_items_limit_caps_expansion(temp_db_path):
-    """Expansion should respect max_context_items limit."""
-    config = AppConfig()
-    config.processing.chunk_size = 32
-    config.search.max_context_items = 2  # Very restrictive
-
-    docling_doc = create_list_document()
-
-    async with HaikuRAG(temp_db_path, config=config, create=True) as client:
-        doc = await create_document_with_docling(client, docling_doc, "Limit Test")
-        assert doc.id is not None
-
-        results = await client.search("grapes", limit=1)
-        assert len(results) > 0
-
-        expanded = await client.expand_context(results)
-
-        # With max_items=2, expansion should be limited
-        content = expanded[0].content.lower()
-        item_count = sum(
-            1 for item in ["apples", "bananas", "oranges", "grapes"] if item in content
-        )
-        # Should have at most 2 items (the limit)
-        assert item_count <= 2, f"Expected at most 2 items, got {item_count}"
-
-
 async def test_expand_context_single_item_document(temp_db_path):
     """Test expand_context with a single-item document."""
     from haiku.rag.store.models.document import Document