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",