From 6945167e1a7334abe4297190352b7b4284ddaee3 Mon Sep 17 00:00:00 2001 From: Tobias Perelstein <5562156+tobocop2@users.noreply.github.com> Date: Fri, 17 Jul 2026 00:48:19 -0400 Subject: [PATCH] Widen retrieved passages with neighboring chunks at context assembly A hit that lands mid-argument loses the text before and after it. With neighbor_expansion set, each selected chunk pulls up to N adjacent chunks per side from its own source and merges them into one contiguous passage, deduplicating the chunker overlap. Expansion runs after the budget fit and spends only leftover tokens, so originals are never dropped for neighbors; citation numbering is untouched and page/line spans are recomputed truthfully. Off by default (neighbor_expansion=0) pending eval evidence. bb-m3mo --- docs/usage.md | 1 + src/lilbee/app/settings_map.py | 6 + src/lilbee/core/config/model.py | 5 + src/lilbee/data/store/core.py | 30 +++- src/lilbee/retrieval/query/neighbors.py | 173 ++++++++++++++++++++++++ src/lilbee/retrieval/query/searcher.py | 56 ++++++-- tests/test_neighbors.py | 156 +++++++++++++++++++++ tests/test_query.py | 68 ++++++++++ tests/test_settings.py | 9 ++ tests/test_store.py | 54 ++++++++ 10 files changed, 548 insertions(+), 10 deletions(-) create mode 100644 src/lilbee/retrieval/query/neighbors.py create mode 100644 tests/test_neighbors.py diff --git a/docs/usage.md b/docs/usage.md index 26724edda..cbbd8550e 100644 --- a/docs/usage.md +++ b/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 | diff --git a/src/lilbee/app/settings_map.py b/src/lilbee/app/settings_map.py index c00d30318..96a07487b 100644 --- a/src/lilbee/app/settings_map.py +++ b/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, diff --git a/src/lilbee/core/config/model.py b/src/lilbee/core/config/model.py index 6f4b0043d..7345c0f82 100644 --- a/src/lilbee/core/config/model.py +++ b/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) diff --git a/src/lilbee/data/store/core.py b/src/lilbee/data/store/core.py index 3bf452922..02cb94514 100644 --- a/src/lilbee/data/store/core.py +++ b/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. diff --git a/src/lilbee/retrieval/query/neighbors.py b/src/lilbee/retrieval/query/neighbors.py new file mode 100644 index 000000000..e874b6785 --- /dev/null +++ b/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), + } + ) diff --git a/src/lilbee/retrieval/query/searcher.py b/src/lilbee/retrieval/query/searcher.py index 93229122a..d1d75e6b8 100644 --- a/src/lilbee/retrieval/query/searcher.py +++ b/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) diff --git a/tests/test_neighbors.py b/tests/test_neighbors.py new file mode 100644 index 000000000..c172cfcdc --- /dev/null +++ b/tests/test_neighbors.py @@ -0,0 +1,156 @@ +"""Tests for neighbor-window expansion (pure logic, mocked store).""" + +from unittest.mock import MagicMock + +from lilbee.data.store import SearchChunk +from lilbee.retrieval.query.neighbors import expand_neighbors, merge_adjacent_texts + + +def _cost(text: str) -> int: + """Mirror Searcher._budget_tokens (3 chars per token, floor 1).""" + return max(1, len(text) // 3) + + +def _chunk( + source="doc.pdf", + index=0, + text="text", + content_type="pdf", + page=1, + line_start=0, + line_end=0, + score=0.9, +) -> SearchChunk: + return SearchChunk( + source=source, + content_type=content_type, + page_start=page, + page_end=page, + line_start=line_start, + line_end=line_end, + chunk=text, + chunk_index=index, + vector=[0.1], + score=score, + ) + + +def _store_with(rows: list[SearchChunk]) -> MagicMock: + """A store mock serving *rows* filtered by the requested source and indices.""" + store = MagicMock() + store.get_chunks_by_indices.side_effect = lambda source, indices: [ + r for r in rows if r.source == source and r.chunk_index in set(indices) + ] + return store + + +class TestMergeAdjacentTexts: + def test_dedups_the_shared_overlap(self): + merged = merge_adjacent_texts(["alpha beta gamma", "gamma delta"]) + assert merged == "alpha beta gamma delta" + + def test_no_overlap_joins_with_a_newline_seam(self): + assert merge_adjacent_texts(["alpha", "delta"]) == "alpha\ndelta" + + def test_fully_contained_text_adds_nothing(self): + assert merge_adjacent_texts(["alpha beta", "beta"]) == "alpha beta" + + def test_merging_an_already_merged_passage_is_idempotent(self): + texts = ["one two", "two three", "three four"] + merged = merge_adjacent_texts(texts) + assert merged == "one two three four" + # Re-merging with either neighbor changes nothing: the neighbor's text + # is fully contained, so a second expansion pass cannot duplicate it. + assert merge_adjacent_texts(["one two", merged]) == merged + assert merge_adjacent_texts([merged, "three four"]) == merged + + +class TestExpandNeighbors: + def test_widens_both_sides_and_recomputes_the_page_span(self): + center = _chunk(index=2, text="gamma delta", page=3) + rows = [ + _chunk(index=1, text="beta gamma", page=2), + _chunk(index=3, text="delta epsilon", page=4), + ] + store = _store_with(rows) + out = expand_neighbors([center], store, radius=1, budget=1000, cost=_cost) + assert len(out) == 1 + assert out[0].chunk == "beta gamma delta epsilon" + assert (out[0].page_start, out[0].page_end) == (2, 4) + assert out[0].score == center.score + assert out[0].chunk_index == center.chunk_index + store.get_chunks_by_indices.assert_called_once_with("doc.pdf", [1, 3]) + + def test_recomputes_the_line_span_for_code(self): + center = _chunk(content_type="code", index=1, text="mid", line_start=10, line_end=20) + rows = [ + _chunk(content_type="code", index=0, text="top", line_start=1, line_end=9), + _chunk(content_type="code", index=2, text="tail", line_start=21, line_end=30), + ] + out = expand_neighbors([center], _store_with(rows), radius=1, budget=1000, cost=_cost) + assert (out[0].line_start, out[0].line_end) == (1, 30) + + def test_single_chunk_document_stays_untouched(self): + center = _chunk(index=0, text="whole doc") + store = _store_with([]) + out = expand_neighbors([center], store, radius=2, budget=1000, cost=_cost) + assert out == [center] + + def test_selected_neighbors_are_never_pulled_twice(self): + # Chunks 2 and 3 are both selected: each widens only away from the + # other, so no passage text is duplicated. + a = _chunk(index=2, text="two three") + b = _chunk(index=3, text="three four") + rows = [ + _chunk(index=1, text="one two"), + _chunk(index=4, text="four five"), + ] + store = _store_with(rows) + out = expand_neighbors([a, b], store, radius=1, budget=1000, cost=_cost) + assert out[0].chunk == "one two three" + assert out[1].chunk == "three four five" + store.get_chunks_by_indices.assert_called_once_with("doc.pdf", [1, 4]) + + def test_higher_ranked_expansion_claims_the_shared_neighbor(self): + first = _chunk(index=2, text="bb") + second = _chunk(index=4, text="dd") + rows = [ + _chunk(index=1, text="aa"), + _chunk(index=3, text="cc"), + _chunk(index=5, text="ee"), + ] + out = expand_neighbors([first, second], _store_with(rows), 1, 1000, _cost) + # Rank order wins: chunk 2 claims index 3, so chunk 4 widens right only. + assert out[0].chunk == "aa\nbb\ncc" + assert out[1].chunk == "dd\nee" + + def test_budget_sheds_farthest_neighbors_first_trailing_on_ties(self): + center = _chunk(index=5, text="e" * 30) + rows = [_chunk(index=i, text=c * 30) for i, c in ((3, "c"), (4, "d"), (6, "f"), (7, "g"))] + # Full widening costs 41 extra tokens; 30 fits only after shedding the + # trailing index 7 (the tie), then the leading index 3 (the farthest). + out = expand_neighbors([center], _store_with(rows), radius=2, budget=30, cost=_cost) + assert out[0].chunk == "\n".join(["d" * 30, "e" * 30, "f" * 30]) + + def test_zero_budget_keeps_every_original(self): + center = _chunk(index=2, text="core") + rows = [_chunk(index=1, text="left"), _chunk(index=3, text="right")] + out = expand_neighbors([center], _store_with(rows), radius=1, budget=0, cost=_cost) + assert out == [center] + + def test_whole_document_selection_expands_nothing(self): + # Every index is a selected passage already, so only the off-the-end + # probe runs and nothing changes. + selected = [_chunk(index=i, text=f"part {i}") for i in range(3)] + store = _store_with([]) + out = expand_neighbors(selected, store, radius=1, budget=1000, cost=_cost) + assert out == selected + store.get_chunks_by_indices.assert_called_once_with("doc.pdf", [3]) + + def test_window_stops_at_a_gap_in_stored_indices(self): + # Index 3 is missing from the store: index 2 must not leapfrog it, + # or the merged passage would splice non-contiguous text. + center = _chunk(index=4, text="dd") + rows = [_chunk(index=2, text="bb"), _chunk(index=5, text="ee")] + out = expand_neighbors([center], _store_with(rows), radius=2, budget=1000, cost=_cost) + assert out[0].chunk == "dd\nee" diff --git a/tests/test_query.py b/tests/test_query.py index ff80799b8..14fb47d78 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -858,6 +858,74 @@ def test_logs_when_trimming(self, mock_svc, caplog): assert "to fit the model context window" in caplog.text +class TestNeighborExpansion: + """Selected chunks widen with adjacent same-source chunks at answer time.""" + + @pytest.fixture(autouse=True) + def _restore(self): + old = (cfg.neighbor_expansion, cfg.num_ctx) + yield + cfg.neighbor_expansion, cfg.num_ctx = old + + def test_off_by_default_never_touches_neighbors(self, mock_svc): + mock_svc.store.search.return_value = [_make_result(chunk="core text")] + mock_svc.provider.chat.return_value = _text_result("answer") + result = get_services().searcher.ask_raw("q") + assert result.sources[0].chunk == "core text" + mock_svc.store.get_chunks_by_indices.assert_not_called() + + def test_widens_the_passage_and_the_prompt(self, mock_svc): + cfg.neighbor_expansion = 1 + mock_svc.store.search.return_value = [_make_result(chunk="core text", chunk_index=2)] + mock_svc.store.get_chunks_by_indices.return_value = [ + _make_result(chunk="before core", chunk_index=1), + _make_result(chunk="text after", chunk_index=3), + ] + mock_svc.provider.chat.return_value = _text_result("answer") + result = get_services().searcher.ask_raw("q") + assert result.sources[0].chunk == "before core text after" + assert result.sources[0].chunk_index == 2 + prompt = mock_svc.provider.chat.call_args[0][0][-1]["content"] + assert "before core text after" in prompt + + def test_widening_keeps_citation_numbering(self, mock_svc): + cfg.neighbor_expansion = 1 + mock_svc.store.search.return_value = [ + _make_result(source="a.pdf", chunk="alpha", chunk_index=5, distance=0.1), + _make_result(source="b.pdf", chunk="beta", chunk_index=0, distance=0.2), + ] + mock_svc.store.get_chunks_by_indices.return_value = [] + mock_svc.provider.chat.return_value = _text_result("From [1] and [2].") + answer = get_services().searcher.ask("q") + assert "1. [a.pdf](file://" in answer + assert "2. [b.pdf](file://" in answer + + def test_sources_block_shows_the_widened_page_span(self, mock_svc): + cfg.neighbor_expansion = 1 + mock_svc.store.search.return_value = [ + _make_result(chunk="core", chunk_index=2, page_start=3, page_end=3) + ] + mock_svc.store.get_chunks_by_indices.return_value = [ + _make_result(chunk="prev", chunk_index=1, page_start=2, page_end=2), + _make_result(chunk="next", chunk_index=3, page_start=4, page_end=4), + ] + mock_svc.provider.chat.return_value = _text_result("answer") + answer = get_services().searcher.ask("q") + assert "pages 2-4" in answer + + def test_tight_budget_sheds_expansion_never_the_original(self, mock_svc): + cfg.neighbor_expansion = 1 + cfg.num_ctx = 1200 + mock_svc.store.search.return_value = [_make_result(chunk="x" * 300, chunk_index=2)] + mock_svc.store.get_chunks_by_indices.return_value = [ + _make_result(chunk="y" * 300, chunk_index=1), + _make_result(chunk="z" * 300, chunk_index=3), + ] + mock_svc.provider.chat.return_value = _text_result("answer") + result = get_services().searcher.ask_raw("q") + assert result.sources[0].chunk == "x" * 300 + + class TestAskRaw: def test_returns_structured_result(self, mock_svc): mock_svc.store.search.return_value = [_make_result(chunk="oil is 5 quarts")] diff --git a/tests/test_settings.py b/tests/test_settings.py index 30f915e51..d41561032 100644 --- a/tests/test_settings.py +++ b/tests/test_settings.py @@ -195,6 +195,15 @@ def test_reranker_fields_in_settings_map(self): assert "reranker_prompt" in SETTINGS_MAP assert SETTINGS_MAP["reranker_type"].choices == ("auto", "cross_encoder", "llm") + def test_neighbor_expansion_in_settings_map(self): + from lilbee.app.settings_map import SETTINGS_MAP, get_default + + defn = SETTINGS_MAP["neighbor_expansion"] + assert defn.writable is True + assert defn.nullable is False + assert defn.group == "Retrieval" + assert get_default("neighbor_expansion") == 0 + class TestReplicaDefaults: """embed/vision replica counts default to 0 = auto (one per GPU at placement).""" diff --git a/tests/test_store.py b/tests/test_store.py index 73514aa36..511362d3a 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -1217,6 +1217,60 @@ def _raise_search(*args, **kwargs): assert len(results) == 1 assert results[0].source == "doc0.md" + +def _one_source_records(source: str, n: int) -> list[dict]: + """*n* sequential chunks all belonging to *source*.""" + records = _make_records(n=n) + for record in records: + record["source"] = source + return records + + +class TestGetChunksByIndices: + def test_returns_requested_indices_in_order(self, store): + store.add_chunks(_one_source_records("a.md", 5)) + results = store.get_chunks_by_indices("a.md", [3, 1]) + assert [r.chunk_index for r in results] == [1, 3] + assert all(r.source == "a.md" for r in results) + + def test_missing_indices_are_absent(self, store): + store.add_chunks(_one_source_records("a.md", 2)) + results = store.get_chunks_by_indices("a.md", [1, 99]) + assert [r.chunk_index for r in results] == [1] + + def test_other_sources_are_excluded(self, store): + store.add_chunks(_one_source_records("a.md", 2) + _one_source_records("b.md", 2)) + results = store.get_chunks_by_indices("a.md", [0, 1]) + assert {r.source for r in results} == {"a.md"} + + def test_empty_indices_returns_empty(self, store): + store.add_chunks(_one_source_records("a.md", 1)) + assert store.get_chunks_by_indices("a.md", []) == [] + + def test_no_table_returns_empty(self, store): + assert store.get_chunks_by_indices("a.md", [0]) == [] + + def test_fallback_when_search_raises(self, store): + """Same Arrow fallback as get_chunks_by_source when table.search() raises.""" + from unittest.mock import patch + + store.add_chunks(_one_source_records("a.md", 3)) + + original_open = store.open_table + + def _broken_open(name): + table = original_open(name) + + def _raise_search(*args, **kwargs): + raise AttributeError("LanceFtsQueryBuilder has no attribute 'metric'") + + table.search = _raise_search + return table + + with patch.object(store, "open_table", side_effect=_broken_open): + results = store.get_chunks_by_indices("a.md", [0, 2]) + assert [r.chunk_index for r in results] == [0, 2] + def test_search_chunk_default_is_raw(self): chunk = SearchChunk( source="a.md",