Widen retrieved passages with neighboring chunks

docs/usage.md

@@ -752,6 +752,7 @@ something feels off.
 | `LILBEE_ADAPTIVE_THRESHOLD_STEP` | `0.2` | How much to widen per step when adaptive threshold triggers |
 | `LILBEE_TEMPORAL_FILTERING` | `true` | When the query contains temporal cues ("recent", "last week"), filter results by document date and sort by recency |
 | `LILBEE_MAX_CONTEXT_SOURCES` | `8` | Max chunks included in the LLM's RAG context. Raise for more coverage, lower for shorter prompts |
+| `LILBEE_NEIGHBOR_EXPANSION` | `0` | Adjacent chunks pulled in on each side of every retrieved passage and merged into one contiguous excerpt, so a hit that lands mid-argument regains its surrounding text. `0` disables |
 | `LILBEE_EXPANSION_SHORT_QUERY_TOKENS` | `2` | Queries this short (in tokens) skip LLM query expansion |
 | `LILBEE_ANN_INDEX_THRESHOLD` | `50000` | Chunk count above which sync builds the ANN vector index |
 | `LILBEE_MAX_REASONING_CHARS` | `64000` | Cap on a reasoning model's thinking output per answer |

src/lilbee/app/settings_map.py

@@ -863,6 +863,12 @@ def get_default(key: str) -> object:
         group=SettingGroup.RETRIEVAL,
         help_text="Maximum unique sources contributing chunks to a single answer",
     ),
+    "neighbor_expansion": SettingDef(
+        int,
+        nullable=False,
+        group=SettingGroup.RETRIEVAL,
+        help_text="Adjacent chunks merged into each retrieved passage per side (0 = off)",
+    ),
     "diversity_max_per_source": SettingDef(
         int,
         nullable=False,

src/lilbee/core/config/model.py

@@ -243,6 +243,11 @@ class Config(BaseSettings):
     # Chunks included in LLM context after adaptive selection.
     max_context_sources: int = ConfigField(default=8, ge=1, writable=True)
 
+    # Adjacent chunks pulled from the same source on each side of every
+    # selected chunk and merged into one contiguous passage, so a hit that
+    # lands mid-argument regains the text before and after it. 0 disables.
+    neighbor_expansion: int = ConfigField(default=0, ge=0, writable=True)
+
     # HyDE (Gao et al. 2022): hypothetical-answer embedding search. +~500ms.
     hyde: bool = ConfigField(default=False, writable=True)
 

src/lilbee/data/store/core.py

@@ -4,7 +4,7 @@
 
 import logging
 import math
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 from contextlib import AbstractContextManager
 from datetime import UTC, datetime
 from pathlib import Path
@@ -859,6 +859,34 @@ def get_chunks_by_source(self, source: str) -> list[SearchChunk]:
             rows = filtered.to_pylist()
         return [SearchChunk(**r) for r in rows]
 
+    def get_chunks_by_indices(self, source: str, indices: Sequence[int]) -> list[SearchChunk]:
+        """Return *source*'s chunks whose ``chunk_index`` is in *indices*.
+
+        Rows come back in ``chunk_index`` order; indices past either end of
+        the document are simply absent from the result.
+        """
+        if not indices:
+            return []
+        table = self.open_table(CHUNKS_TABLE)
+        if table is None:
+            return []
+        escaped = escape_sql_string(source)
+        wanted = ", ".join(str(int(i)) for i in indices)
+        try:
+            predicate = f"source = '{escaped}' AND chunk_index IN ({wanted})"
+            rows = table.search().where(predicate).limit(None).to_list()
+        except Exception:
+            # Same FTS-enabled query-builder limitation as get_chunks_by_source;
+            # fall through to a pyarrow.compute filter on the Arrow table.
+            log.debug("get_chunks_by_indices search() failed, using Arrow fallback", exc_info=True)
+            arrow_tbl = table.to_arrow()
+            mask = pc.and_(
+                pc.equal(arrow_tbl["source"], source),
+                pc.is_in(arrow_tbl["chunk_index"], value_set=pa.array(list(indices), pa.int32())),
+            )
+            rows = arrow_tbl.filter(mask).to_pylist()
+        return sorted((SearchChunk(**r) for r in rows), key=lambda c: c.chunk_index)
+
     def _delete_by_sources_unlocked(self, sources: list[str]) -> None:
         """Delete the sources' chunks, page texts, and chunk-concept rows.
 

src/lilbee/retrieval/query/neighbors.py

@@ -0,0 +1,173 @@
+"""Neighbor-window expansion: widen retrieved passages with adjacent chunks.
+
+A hit that lands in the middle of an argument loses the sentences before and
+after it. After context selection, each selected chunk pulls up to N adjacent
+chunks per side from its own source and merges them into one contiguous
+passage, deduplicating the overlap text adjacent chunks share from chunking.
+The widened passage keeps the original chunk's score and citation identity;
+only its text and page/line span change.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from lilbee.data.store import SearchChunk, Store
+
+
+def _overlap_chars(left: str, right: str) -> int:
+    """Length of the longest suffix of *left* that is a prefix of *right*.
+
+    Adjacent chunks carry the chunker's overlap verbatim (both cuts come from
+    the same document text), so the longest match IS the shared region. A
+    fully contained text matches whole, which is what makes merging an
+    already-widened passage idempotent instead of duplicating its neighbors.
+    """
+    for k in range(min(len(left), len(right)), 0, -1):
+        if left.endswith(right[:k]):
+            return k
+    return 0
+
+
+def merge_adjacent_texts(texts: list[str]) -> str:
+    """Concatenate adjacent chunk texts, deduplicating their shared overlap.
+
+    Adjacent texts with no detectable overlap (a chunk_overlap=0 build) join
+    with a newline seam rather than gluing two words together.
+    """
+    merged = texts[0]
+    for text in texts[1:]:
+        k = _overlap_chars(merged, text)
+        tail = text[k:]
+        if not tail:
+            continue
+        merged = merged + tail if k else f"{merged}\n{tail}"
+    return merged
+
+
+def expand_neighbors(
+    results: list[SearchChunk],
+    store: Store,
+    radius: int,
+    budget: int,
+    cost: Callable[[str], int],
+) -> list[SearchChunk]:
+    """Widen each result with up to *radius* adjacent same-source chunks.
+
+    Results are processed in rank order and only spend *budget*, the tokens
+    left over after the originals were fitted: a widened passage whose extra
+    cost does not fit sheds its farthest neighbors first and falls back to
+    the original text, so expansion is always trimmed before any original
+    chunk. An index that is itself selected, or already claimed by a
+    higher-ranked expansion, is never pulled again, so no passage text is
+    duplicated (a document routed whole expands to nothing).
+    """
+    centers: dict[str, set[int]] = {}
+    for r in results:
+        centers.setdefault(r.source, set()).add(r.chunk_index)
+    rows = _fetch_neighbor_rows(store, centers, radius)
+    if not rows:
+        return results
+    claimed = {source: set(indices) for source, indices in centers.items()}
+    remaining = budget
+    expanded: list[SearchChunk] = []
+    for r in results:
+        widened, spent = _widen(r, rows, claimed[r.source], radius, remaining, cost)
+        remaining -= spent
+        expanded.append(widened)
+    return expanded
+
+
+def _fetch_neighbor_rows(
+    store: Store, centers: dict[str, set[int]], radius: int
+) -> dict[tuple[str, int], SearchChunk]:
+    """Every candidate neighbor row, fetched with one store call per source.
+
+    Selected indices are excluded up front (they are already passages);
+    indices past the end of a document are simply absent from the reply.
+    """
+    rows: dict[tuple[str, int], SearchChunk] = {}
+    for source, owned in centers.items():
+        wanted = sorted(
+            {
+                index
+                for center in owned
+                for index in range(center - radius, center + radius + 1)
+                if index >= 0 and index not in owned
+            }
+        )
+        for row in store.get_chunks_by_indices(source, wanted):
+            rows[(source, row.chunk_index)] = row
+    return rows
+
+
+def _neighbor_run(
+    result: SearchChunk,
+    rows: dict[tuple[str, int], SearchChunk],
+    claimed: set[int],
+    step: int,
+    radius: int,
+) -> list[int]:
+    """Contiguous free neighbor indices on one side of the center, nearest first."""
+    indices: list[int] = []
+    for offset in range(1, radius + 1):
+        index = result.chunk_index + step * offset
+        if index in claimed or (result.source, index) not in rows:
+            break
+        indices.append(index)
+    return indices
+
+
+def _widen(
+    result: SearchChunk,
+    rows: dict[tuple[str, int], SearchChunk],
+    claimed: set[int],
+    radius: int,
+    remaining: int,
+    cost: Callable[[str], int],
+) -> tuple[SearchChunk, int]:
+    """One result widened within *remaining* tokens: (chunk, tokens spent).
+
+    Neighbors extend from the center until a missing, selected, or already
+    claimed index stops each side. While the widened text over-spends, the
+    farthest neighbor is shed first (a tie sheds the trailing side, keeping
+    the text that leads up to the hit); shedding everything keeps the
+    original chunk untouched.
+    """
+    center = result.chunk_index
+    left = _neighbor_run(result, rows, claimed, -1, radius)
+    right = _neighbor_run(result, rows, claimed, +1, radius)
+    while left or right:
+        span = sorted([*left, center, *right])
+        texts = [
+            result.chunk if index == center else rows[(result.source, index)].chunk
+            for index in span
+        ]
+        merged = merge_adjacent_texts(texts)
+        extra = cost(merged) - cost(result.chunk)
+        if extra <= remaining:
+            claimed.update(index for index in span if index != center)
+            neighbors = [rows[(result.source, index)] for index in span if index != center]
+            return _widened_copy(result, merged, neighbors), extra
+        if right and (not left or right[-1] - center >= center - left[-1]):
+            right.pop()
+        else:
+            left.pop()
+    return result, 0
+
+
+def _widened_copy(result: SearchChunk, merged: str, neighbors: list[SearchChunk]) -> SearchChunk:
+    """The result with widened text and a truthfully recomputed page/line span."""
+    spans = [result, *neighbors]
+    return result.model_copy(
+        update={
+            "chunk": merged,
+            "page_start": min(s.page_start for s in spans),
+            "page_end": max(s.page_end for s in spans),
+            "line_start": min(s.line_start for s in spans),
+            "line_end": max(s.line_end for s in spans),
+        }
+    )

src/lilbee/retrieval/query/searcher.py

@@ -61,6 +61,7 @@
     title_candidates,
 )
 from lilbee.retrieval.query.memory import format_memory_block
+from lilbee.retrieval.query.neighbors import expand_neighbors
 from lilbee.retrieval.query.tokenize import _idf_weights, _tokenize
 from lilbee.retrieval.reasoning import (
     StreamToken,
@@ -797,6 +798,7 @@ def _finalize_context(
         """
         system = self._system_with_memory(self._config.rag_system_prompt, question)
         results = self._fit_context_budget(results, system, question, history, scale)
+        results = self._widen_with_neighbors(results, system, question, history, scale)
         context = build_context(results)
         prompt = CONTEXT_TEMPLATE.format(context=context, question=question)
         messages: list[ChatMessage] = [{"role": "system", "content": system}]
@@ -810,19 +812,14 @@ def _budget_tokens(text: str) -> int:
         """Conservative token cost for budgeting (see _BUDGET_CHARS_PER_TOKEN)."""
         return max(1, len(text) // _BUDGET_CHARS_PER_TOKEN)
 
-    def _fit_context_budget(
+    def _context_budget(
         self,
-        results: list[SearchChunk],
         system: str,
         question: str,
         history: list[ChatMessage] | None,
         scale: float = 1.0,
-    ) -> list[SearchChunk]:
-        """Drop the lowest-ranked sources until the assembled prompt fits num_ctx.
-
-        ``max_context_sources`` caps by count; this caps by tokens so a
-        retrieval-heavy query degrades gracefully instead of erroring with
-        CONTEXT_OVERFLOW. The top-ranked source is always kept.
+    ) -> int:
+        """Token budget left for source passages after the fixed prompt parts.
 
         The ceiling is the engine's ACTUAL per-slot window when known: the
         configured value is a target the dynamic picker aims for, and the
@@ -840,7 +837,23 @@ def _fit_context_budget(
             + _ANSWER_RESERVE_TOKENS
             + _CONTEXT_TEMPLATE_TOKENS
         )
-        budget = int((ctx - reserve) * scale)
+        return int((ctx - reserve) * scale)
+
+    def _fit_context_budget(
+        self,
+        results: list[SearchChunk],
+        system: str,
+        question: str,
+        history: list[ChatMessage] | None,
+        scale: float = 1.0,
+    ) -> list[SearchChunk]:
+        """Drop the lowest-ranked sources until the assembled prompt fits num_ctx.
+
+        ``max_context_sources`` caps by count; this caps by tokens so a
+        retrieval-heavy query degrades gracefully instead of erroring with
+        CONTEXT_OVERFLOW. The top-ranked source is always kept.
+        """
+        budget = self._context_budget(system, question, history, scale)
         kept: list[SearchChunk] = []
         used = 0
         for r in results:
@@ -857,6 +870,31 @@ def _fit_context_budget(
             )
         return kept
 
+    def _widen_with_neighbors(
+        self,
+        results: list[SearchChunk],
+        system: str,
+        question: str,
+        history: list[ChatMessage] | None,
+        scale: float,
+    ) -> list[SearchChunk]:
+        """Widen each fitted passage with adjacent same-source chunks.
+
+        Runs after the budget fit so neighbors only ever spend leftover
+        budget: a tight window sheds expansion first and never drops an
+        original chunk for a neighbor. Widening changes a passage's text
+        and page/line span only, so citation numbering and the sources
+        block are untouched.
+        """
+        radius = self._config.neighbor_expansion
+        if radius <= 0:
+            return results
+        budget = self._context_budget(system, question, history, scale)
+        used = sum(self._budget_tokens(r.chunk) + _PER_SOURCE_TOKENS for r in results)
+        return expand_neighbors(
+            results, self._store, radius, max(0, budget - used), self._budget_tokens
+        )
+
     def _system_with_memory(self, base_prompt: str, question: str) -> str:
         """Append the local-owner memory block to *base_prompt* when memory is enabled."""
         block = self._memory_block(question)