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 01/28] 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", From 046244ec6333428c5eda6030c4f55ef68a0990a0 Mon Sep 17 00:00:00 2001 From: Tobias Perelstein <5562156+tobocop2@users.noreply.github.com> Date: Fri, 17 Jul 2026 01:09:07 -0400 Subject: [PATCH 02/28] Make document titles searchable as a weighted lexical arm The lexical index covered only chunk text, so BM25 could never match a document's title or filename. Every chunk now carries a nullable title (extraction metadata when available, cleaned filename stem otherwise) with its own FTS index, queried as a third RRF arm at a configurable weight when title_search is on (off by default pending eval evidence). Existing chunk-FTS call sites pin fts_columns explicitly: once a second FTS index exists, an unpinned query searches every indexed column, which would silently widen BM25 semantics. Old indexes keep working: schemas evolve in place and the title arm contributes nothing until reindex. Source rows also persist extraction title/authors/created_at for future filtering. bb-t8pr --- docs/usage.md | 2 + src/lilbee/app/settings_map.py | 12 ++ src/lilbee/core/config/model.py | 9 + src/lilbee/data/export.py | 9 +- src/lilbee/data/ingest/extract.py | 11 +- src/lilbee/data/ingest/pipeline.py | 20 ++- src/lilbee/data/ingest/title.py | 38 +++++ src/lilbee/data/ingest/types.py | 10 +- src/lilbee/data/store/__init__.py | 2 + src/lilbee/data/store/core.py | 111 +++++++++++-- src/lilbee/data/store/fusion.py | 70 +++++--- src/lilbee/data/store/lance_helpers.py | 6 +- src/lilbee/data/store/schema.py | 3 + src/lilbee/data/store/types.py | 23 +++ tests/server/test_handlers.py | 1 + tests/test_export.py | 9 + tests/test_formats.py | 1 + tests/test_ingest.py | 60 ++++++- tests/test_ingest_title.py | 57 +++++++ tests/test_settings.py | 28 ++++ tests/test_store.py | 218 ++++++++++++++++++++++++- tests/test_store_fusion.py | 56 +++++++ 22 files changed, 706 insertions(+), 50 deletions(-) create mode 100644 src/lilbee/data/ingest/title.py create mode 100644 tests/test_ingest_title.py diff --git a/docs/usage.md b/docs/usage.md index 26724edda..8b2d2a2f6 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -748,6 +748,8 @@ something feels off. | `LILBEE_RERANK_BLEND` | `true` | Blend reranker scores with retrieval fusion; off = the reranker's own ordering stands | | `LILBEE_HYDE` | `false` | Enable Hypothetical Document Embeddings: an LLM drafts a hypothetical answer, that's embedded, and results are merged with the original query's. Adds ~500 ms per query; helps on vague questions | | `LILBEE_HYDE_WEIGHT` | `0.7` | How much to trust HyDE results relative to the direct query (0.0-1.0) | +| `LILBEE_TITLE_SEARCH` | `false` | Match queries against document titles as a third hybrid-search arm, so a query naming a document by title surfaces its chunks | +| `LILBEE_TITLE_SEARCH_WEIGHT` | `0.5` | Title arm weight in rank fusion (1.0 = equal voice with the vector and text arms) | | `LILBEE_ADAPTIVE_THRESHOLD` | `false` | Widen the distance threshold step by step when too few results pass, on the vector-only fallback path (no effect on the default hybrid search) | | `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 | diff --git a/src/lilbee/app/settings_map.py b/src/lilbee/app/settings_map.py index c00d30318..43f1c4ee8 100644 --- a/src/lilbee/app/settings_map.py +++ b/src/lilbee/app/settings_map.py @@ -817,6 +817,18 @@ def get_default(key: str) -> object: group=SettingGroup.RETRIEVAL, help_text="Candidate-pool multiplier over top_k before reranking", ), + "title_search": SettingDef( + bool, + nullable=False, + group=SettingGroup.RETRIEVAL, + help_text="Match queries against document titles as a third hybrid-search arm", + ), + "title_search_weight": SettingDef( + float, + nullable=False, + group=SettingGroup.RETRIEVAL, + help_text="Title arm weight in rank fusion (1.0 = equal voice with the other arms)", + ), "history_rewrite": SettingDef( bool, nullable=False, diff --git a/src/lilbee/core/config/model.py b/src/lilbee/core/config/model.py index 6f4b0043d..642473975 100644 --- a/src/lilbee/core/config/model.py +++ b/src/lilbee/core/config/model.py @@ -194,6 +194,15 @@ class Config(BaseSettings): # fusion arms stay exactly top_k deep. candidate_multiplier: int = ConfigField(default=3, ge=1, writable=True) + # Third lexical arm in hybrid search: BM25 over document titles, fused with + # the vector and chunk arms so a query naming a document by title surfaces + # its chunks. Off by default until the eval harness measures it. + title_search: bool = ConfigField(default=False, writable=True) + + # Title arm weight relative to a full arm in rank fusion (1.0 = equal voice + # with the vector and chunk arms). + title_search_weight: float = ConfigField(default=0.5, ge=0.0, le=1.0, writable=True) + # Chunk count at/above which sync builds an approximate (ANN) vector index # so search stays fast at millions of vectors. Below this, search uses exact # flat scan (faster and exact for small vaults). 0 disables the ANN index. diff --git a/src/lilbee/data/export.py b/src/lilbee/data/export.py index 3ad1e080c..5a2703374 100644 --- a/src/lilbee/data/export.py +++ b/src/lilbee/data/export.py @@ -10,7 +10,8 @@ from typing import TYPE_CHECKING, cast from lilbee.data.ingest.extract import chunk_and_embed_pages -from lilbee.data.store import ChunkWrite, PageTextRecord, SourceType +from lilbee.data.ingest.title import derive_title +from lilbee.data.store import ChunkWrite, PageTextRecord, SourceMeta, SourceType from lilbee.runtime.progress import DetailedProgressCallback, noop_callback if TYPE_CHECKING: @@ -225,6 +226,11 @@ async def import_dataset( content_type = source_rows[0]["content_type"] or "text" page_texts = [(r["page"], r["text"]) for r in source_rows] chunks = await chunk_and_embed_pages(page_texts, name, content_type, on_progress) + # Imports carry no extraction metadata; the stem-derived title keeps + # imported chunks visible to the title search arm. + title = derive_title(name) + for chunk in chunks: + chunk["title"] = title # One locked transaction (cleanup + chunks + page texts + source row) so a # failure can't leave the source with its old rows deleted and no new ones; # the embedding-dim check inside runs before the cleanup delete. @@ -238,6 +244,7 @@ async def import_dataset( needs_cleanup=True, page_texts=[dict(r) for r in source_rows], source_type=SourceType.IMPORTED, + meta=SourceMeta(title=title), ) ], ) diff --git a/src/lilbee/data/ingest/extract.py b/src/lilbee/data/ingest/extract.py index 0aea58d4f..4f56fdabb 100644 --- a/src/lilbee/data/ingest/extract.py +++ b/src/lilbee/data/ingest/extract.py @@ -22,6 +22,7 @@ from lilbee.data.ingest.discovery import file_hash from lilbee.data.ingest.ocr_cache import load_ocr_pages, ocr_cache_key, store_ocr_pages from lilbee.data.ingest.offload import to_ingest_thread +from lilbee.data.ingest.title import source_meta_from_extraction from lilbee.data.ingest.types import ( IMAGE_CONTENT_TYPE, MARKDOWN_OUTPUT, @@ -31,7 +32,7 @@ ChunkRecord, ExtractMode, ) -from lilbee.data.store import ChunkType, PageTextRecord +from lilbee.data.store import ChunkType, PageTextRecord, SourceMeta from lilbee.runtime.cancellation import TaskCancelledError from lilbee.runtime.cpu import cpu_quota from lilbee.runtime.progress import ( @@ -549,11 +550,14 @@ async def ingest_document( quiet: bool = False, on_progress: DetailedProgressCallback = noop_callback, page_texts_out: list[PageTextRecord] | None = None, + meta_out: list[SourceMeta] | None = None, ) -> list[ChunkRecord]: """Extract and chunk a document, embed, return records. Vision OCR is controlled by ``cfg.enable_ocr`` (see ``_should_run_ocr``). When ``page_texts_out`` is given, per-page text is appended for export. + When ``meta_out`` is given, the document's extraction metadata (title, + authors, creation date) is appended for the source row. """ # An image carries no text layer; route it straight to OCR (vision or Tesseract) # instead of a no-op kreuzberg markdown extract that yields nothing for a scan. @@ -567,6 +571,11 @@ async def ingest_document( config = extraction_config(content_type_to_mode(content_type)) result = await to_ingest_thread(extract_file_sync, str(path), config=config) + # Captured before the scanned-PDF fallback: a scan's PDF metadata (title, + # authors) survives even when its text layer is empty. + if meta_out is not None: + meta_out.append(source_meta_from_extraction(result.metadata or {}, source_name)) + if content_type == PDF_CONTENT_TYPE and not _has_meaningful_text(result): return await _handle_scanned_pdf_fallback( path, diff --git a/src/lilbee/data/ingest/pipeline.py b/src/lilbee/data/ingest/pipeline.py index 1173cb6ee..0a1833d9a 100644 --- a/src/lilbee/data/ingest/pipeline.py +++ b/src/lilbee/data/ingest/pipeline.py @@ -41,6 +41,7 @@ write_skip_markers, write_skip_reasons, ) +from lilbee.data.ingest.title import derive_title from lilbee.data.ingest.types import ( ChunkRecord, FileChangePlan, @@ -53,6 +54,7 @@ ChunkWrite, ConceptRecords, PageTextRecord, + SourceMeta, SourceRecord, SourceStat, SourceStatBackfill, @@ -198,21 +200,26 @@ async def _produce_records( quiet: bool = False, on_progress: DetailedProgressCallback = noop_callback, page_texts_out: list[PageTextRecord] | None = None, + meta_out: list[SourceMeta] | None = None, ) -> list[ChunkRecord]: """Extract, chunk, and embed a single file into store-ready records. The LanceDB write is deferred: records are returned to the caller and written in a batched flush (see :func:`_flush_writes`), so bulk ingest pays one write-lock acquisition per batch instead of one per file. The per-page text - dataset rows land in ``page_texts_out`` and are written by the same flush. + dataset rows land in ``page_texts_out`` and are written by the same flush; + the document's metadata (extraction-provided when available, stem-derived + title otherwise) lands in ``meta_out`` and stamps every record's ``title``. """ records: list[ChunkRecord] page_texts: list[PageTextRecord] = page_texts_out if page_texts_out is not None else [] + meta = SourceMeta(title=derive_title(source_name)) if content_type == "code": records = await to_ingest_thread(ingest_code_sync, path, source_name, on_progress) elif path.suffix.lower() == ".md": records = await ingest_markdown(path, source_name, on_progress, page_texts_out=page_texts) else: + extracted_meta: list[SourceMeta] = [] records = await ingest_document( path, source_name, @@ -220,8 +227,15 @@ async def _produce_records( quiet=quiet, on_progress=on_progress, page_texts_out=page_texts, + meta_out=extracted_meta, ) + if extracted_meta: + meta = extracted_meta[0] + for record in records: + record["title"] = meta.title + if meta_out is not None: + meta_out.append(meta) return records @@ -680,6 +694,7 @@ async def _process_one(entry: FileToProcess, file_index: int) -> _IngestResult: # transaction as the new write (see _flush_writes), so cleanup is # carried on the result rather than run eagerly here. page_texts: list[PageTextRecord] = [] + meta_out: list[SourceMeta] = [] records = await _produce_records( entry.path, name, @@ -687,6 +702,7 @@ async def _process_one(entry: FileToProcess, file_index: int) -> _IngestResult: quiet=quiet, on_progress=on_progress, page_texts_out=page_texts, + meta_out=meta_out, ) concept_records = await _build_concept_records(records, name) entity_rows = await _build_entity_records(records, name) @@ -707,6 +723,7 @@ async def _process_one(entry: FileToProcess, file_index: int) -> _IngestResult: stat=entry.stat, concept_records=concept_records, entity_rows=entity_rows, + meta=meta_out[0] if meta_out else None, ) except (asyncio.CancelledError, TaskCancelledError) as exc: # TaskCancelledError is the TUI's cooperative cancel signal raised @@ -1031,6 +1048,7 @@ def _flush_batch(buffer: list[_IngestResult]) -> None: needs_cleanup=r.needs_cleanup, stat=r.stat, page_texts=cast(list[dict], r.page_texts or []), + meta=r.meta, ) for r in buffer ] diff --git a/src/lilbee/data/ingest/title.py b/src/lilbee/data/ingest/title.py new file mode 100644 index 000000000..c34003002 --- /dev/null +++ b/src/lilbee/data/ingest/title.py @@ -0,0 +1,38 @@ +"""Document title and source-level metadata derivation for ingest.""" + +from __future__ import annotations + +import re +from collections.abc import Mapping +from pathlib import PurePath +from typing import Any + +from lilbee.data.store import SourceMeta + +# Filename-stem separators flattened to spaces when no extracted title exists. +_STEM_SEPARATOR_RE = re.compile(r"[_\-\s]+") + + +def derive_title(source_name: str, metadata_title: str | None = None) -> str: + """Human-readable document title: the extracted title, else the cleaned filename stem. + + The stem cleanup flattens underscore/hyphen separators to spaces so BM25 + tokenizes ``survey_214.pdf`` into the same terms a query would use. + """ + if metadata_title and metadata_title.strip(): + return metadata_title.strip() + return _STEM_SEPARATOR_RE.sub(" ", PurePath(source_name).stem).strip() + + +def source_meta_from_extraction(metadata: Mapping[str, Any], source_name: str) -> SourceMeta: + """Fold kreuzberg extraction metadata into a :class:`SourceMeta`. + + The title falls back to the filename stem; authors and creation date stay + empty (persisted NULL) when the extractor reports none. + """ + authors = metadata.get("authors") or [] + return SourceMeta( + title=derive_title(source_name, metadata.get("title")), + authors=", ".join(a for a in authors if a), + created_at=str(metadata.get("created_at") or ""), + ) diff --git a/src/lilbee/data/ingest/types.py b/src/lilbee/data/ingest/types.py index 111269790..8ea4da8b1 100644 --- a/src/lilbee/data/ingest/types.py +++ b/src/lilbee/data/ingest/types.py @@ -5,7 +5,7 @@ from dataclasses import dataclass from enum import StrEnum from pathlib import Path -from typing import NamedTuple, TypedDict +from typing import NamedTuple, NotRequired, TypedDict from pydantic import BaseModel @@ -13,6 +13,7 @@ ChunkType, ConceptRecords, PageTextRecord, + SourceMeta, SourceStat, SourceStatBackfill, ) @@ -72,6 +73,9 @@ class ChunkRecord(TypedDict): chunk: str chunk_index: int vector: list[float] + # Stamped once per document by the pipeline (see _produce_records) rather + # than set by every producer, so the derivation lives in one place. + title: NotRequired[str] class SyncResult(BaseModel): @@ -124,7 +128,8 @@ class _IngestResult: travels with the records so the flush can delete the source's old chunks in the same transaction. ``page_texts`` carries the per-page text dataset rows and ``concept_records`` the file's concept-table rows, and ``entity_rows`` - the file's typed-entity rows, all written by the same flush. + the file's typed-entity rows, all written by the same flush. ``meta`` + carries the document's extraction-time metadata for the source row. """ name: str @@ -138,6 +143,7 @@ class _IngestResult: stat: SourceStat | None = None concept_records: ConceptRecords | None = None entity_rows: list[dict] | None = None + meta: SourceMeta | None = None # Extension → content_type string for document formats handled by kreuzberg. diff --git a/src/lilbee/data/store/__init__.py b/src/lilbee/data/store/__init__.py index ab2b4f3fb..83a42f05f 100644 --- a/src/lilbee/data/store/__init__.py +++ b/src/lilbee/data/store/__init__.py @@ -28,6 +28,7 @@ RemoveResult, SearchChunk, SearchScope, + SourceMeta, SourceRecord, SourceStat, SourceStatBackfill, @@ -53,6 +54,7 @@ "RemoveResult", "SearchChunk", "SearchScope", + "SourceMeta", "SourceRecord", "SourceStat", "SourceStatBackfill", diff --git a/src/lilbee/data/store/core.py b/src/lilbee/data/store/core.py index 3bf452922..56e1bcefd 100644 --- a/src/lilbee/data/store/core.py +++ b/src/lilbee/data/store/core.py @@ -64,6 +64,7 @@ PageTextRecord, RemoveResult, SearchChunk, + SourceMeta, SourceRecord, SourceStat, SourceStatBackfill, @@ -118,6 +119,14 @@ def _drop_unsupported_far_rows( # in place with the SOURCE_STAT_UNKNOWN sentinel. _SOURCE_STAT_COLUMNS = ("size_bytes", "mtime_ns", "stat_captured_ns") +# Extraction-metadata columns of ``_sources``; nullable strings, so legacy +# tables migrate in place with NULL (meaning "extractor reported nothing"). +_SOURCE_META_COLUMNS = ("title", "authors", "created_at") + +# Document-title column of the chunks table; nullable so pre-title rows and +# writers that carry no title (wiki pages) read as NULL. +_TITLE_COLUMN = "title" + # (table, source column) pairs deleted when a source's rows are replaced. The # concept nodes/edges tables carry no source column (corpus-level aggregates), # so only the per-chunk concept mapping is source-scoped. @@ -218,10 +227,18 @@ def _chunks_schema(self) -> pa.Schema: pa.field("line_end", pa.int32()), pa.field("chunk", pa.utf8()), pa.field("chunk_index", pa.int32()), + pa.field(_TITLE_COLUMN, pa.utf8()), pa.field("vector", pa.list_(pa.float32(), self._config.embedding_dim)), ] ) + def _chunks_table(self) -> lancedb.table.Table: + """Open/create the chunks table, adding the title column to pre-title tables.""" + table = ensure_table(self.get_db(), CHUNKS_TABLE, self._chunks_schema()) + if _TITLE_COLUMN not in table.schema.names: + table.add_columns({_TITLE_COLUMN: "CAST(NULL AS STRING)"}) + return table + def get_meta(self) -> StoreMeta | None: """Return the persisted store metadata row, or ``None`` if unset.""" table = self.open_table(META_TABLE) @@ -410,15 +427,32 @@ def ensure_fts_index(self) -> None: return try: if _has_fts_index(table): + # One optimize folds new rows into every index on the + # table, the title index included. table.optimize() log.debug("FTS index optimized on '%s'", CHUNKS_TABLE) else: table.create_fts_index("chunk", replace=False) log.debug("FTS index created on '%s'", CHUNKS_TABLE) + self._ensure_title_fts_unlocked(table) self._fts_ready = True except Exception: log.debug("FTS index ensure failed (empty table?)", exc_info=True) + def _ensure_title_fts_unlocked(self, table: lancedb.table.Table) -> None: + """Create the title FTS index when the column exists. Caller holds ``write_lock()``. + + Failure never blocks the chunk index: the title arm feature-detects the + index per query, so a store without it simply searches without titles. + """ + if _TITLE_COLUMN not in table.schema.names or _has_fts_index(table, _TITLE_COLUMN): + return + try: + table.create_fts_index(_TITLE_COLUMN, replace=False) + log.debug("Title FTS index created on '%s'", CHUNKS_TABLE) + except Exception: + log.debug("Title FTS index create failed", exc_info=True) + def ensure_vector_index(self, *, force: bool = False) -> bool: """Build or refresh the ANN vector index when the corpus is large enough. @@ -473,8 +507,7 @@ def add_chunks(self, records: list[dict]) -> int: if not records: return 0 _check_vector_dims(records, embedding_dim) - db = self.get_db() - table = ensure_table(db, CHUNKS_TABLE, self._chunks_schema()) + table = self._chunks_table() table.add(records) if self.get_meta() is None: self._write_meta_unlocked( @@ -497,7 +530,10 @@ def bm25_probe( if not self._fts_ready: return [] try: - query = table.search(query_text, query_type="fts") + # Pin the chunk column: once a title FTS index exists, an unpinned + # FTS query searches every indexed column, which would silently + # widen the probe's semantics. + query = table.search(query_text, query_type="fts", fts_columns="chunk") if chunk_type: query = query.where(_chunk_type_predicate(chunk_type)) rows = query.limit(top_k).to_list() @@ -600,8 +636,32 @@ def _fts_arm( limit: int, chunk_type: ChunkType | None, ) -> list[SearchChunk]: - """BM25-arm candidates.""" - query = table.search(query_text, query_type="fts").limit(limit) + """BM25-arm candidates over the chunk text. + + The column is pinned: once a title FTS index exists, an unpinned FTS + query searches every indexed column, which would double-count title + tokens this arm's sibling already scores. + """ + query = table.search(query_text, query_type="fts", fts_columns="chunk").limit(limit) + if chunk_type: + query = query.where(_chunk_type_predicate(chunk_type)) + return [SearchChunk(**r) for r in query.to_list()] + + def _title_arm( + self, + table: lancedb.table.Table, + query_text: str, + limit: int, + chunk_type: ChunkType | None, + ) -> list[SearchChunk]: + """BM25-arm candidates over the document title column. + + Empty when the store predates the title column or its FTS index: old + indexes keep working, the arm simply contributes nothing there. + """ + if not _has_fts_index(table, _TITLE_COLUMN): + return [] + query = table.search(query_text, query_type="fts", fts_columns=_TITLE_COLUMN).limit(limit) if chunk_type: query = query.where(_chunk_type_predicate(chunk_type)) return [SearchChunk(**r) for r in query.to_list()] @@ -628,10 +688,20 @@ def _hybrid_search( identifiers), which is exactly what MMR penalizes: selecting the fused pool through MMR measurably traded relevant lexical hits for diverse off-topic neighbors on graded evaluation. + + When ``cfg.title_search`` is on, a third BM25 arm over document + titles joins the fusion at ``cfg.title_search_weight``; title rows + carry ``bm25_score``, so a title match counts as lexical support + for the distance exemption like any other lexical hit. """ + title_rows: list[SearchChunk] = [] + if self._config.title_search: + title_rows = self._title_arm(table, query_text, top_k, chunk_type) fused = fuse_arms( self._vector_arm(table, query_vector, top_k, chunk_type), self._fts_arm(table, query_text, top_k, chunk_type), + title_rows, + title_weight=self._config.title_search_weight, ) fused = _drop_unsupported_far_rows(fused, max_distance) return fused[:top_k] @@ -964,8 +1034,14 @@ def _source_row( chunk_count: int, source_type: str, stat: SourceStat | None, + meta: SourceMeta | None = None, ) -> dict: - """Build one ``_sources`` row, defaulting absent stat to the unknown sentinel.""" + """Build one ``_sources`` row, defaulting absent stat to the unknown sentinel. + + Absent extraction metadata persists as NULL, matching rows written + before the metadata columns existed. + """ + meta = meta or SourceMeta() return { "filename": filename, "file_hash": file_hash, @@ -975,14 +1051,19 @@ def _source_row( "size_bytes": stat.size_bytes if stat else SOURCE_STAT_UNKNOWN, "mtime_ns": stat.mtime_ns if stat else SOURCE_STAT_UNKNOWN, "stat_captured_ns": stat.captured_ns if stat else SOURCE_STAT_UNKNOWN, + "title": meta.title or None, + "authors": meta.authors or None, + "created_at": meta.created_at or None, } def _sources_table(self) -> lancedb.table.Table: - """Open/create ``_sources``, adding the stat columns to pre-stat tables.""" + """Open/create ``_sources``, adding the stat and metadata columns to older tables.""" table = ensure_table(self.get_db(), SOURCES_TABLE, _sources_schema()) - missing = [name for name in _SOURCE_STAT_COLUMNS if name not in table.schema.names] + defaults = {name: f"CAST({SOURCE_STAT_UNKNOWN} AS BIGINT)" for name in _SOURCE_STAT_COLUMNS} + defaults |= {name: "CAST(NULL AS STRING)" for name in _SOURCE_META_COLUMNS} + missing = {name: sql for name, sql in defaults.items() if name not in table.schema.names} if missing: - table.add_columns({name: f"CAST({SOURCE_STAT_UNKNOWN} AS BIGINT)" for name in missing}) + table.add_columns(missing) return table def _replace_source_rows_unlocked(self, rows: list[dict]) -> None: @@ -1007,9 +1088,10 @@ def upsert_source( chunk_count: int, source_type: SourceType = SourceType.DOCUMENT, stat: SourceStat | None = None, + meta: SourceMeta | None = None, ) -> None: """Add or update a source tracking record.""" - row = self._source_row(filename, file_hash, chunk_count, source_type, stat) + row = self._source_row(filename, file_hash, chunk_count, source_type, stat, meta) with self._write_lock(): self._replace_source_rows_unlocked([row]) self._invalidate_source_cache() @@ -1067,7 +1149,7 @@ def write_chunks_batch(self, items: list[ChunkWrite]) -> int: db = self.get_db() self._cleanup_batch_unlocked(items) self._add_page_texts_unlocked(db, items) - self._add_chunk_records_unlocked(db, all_records, embedding_model, embedding_dim) + self._add_chunk_records_unlocked(all_records, embedding_model, embedding_dim) self._replace_source_rows_unlocked(self._batch_source_rows(items)) self._invalidate_source_cache() return len(all_records) @@ -1086,7 +1168,6 @@ def _add_page_texts_unlocked(self, db: lancedb.DBConnection, items: list[ChunkWr def _add_chunk_records_unlocked( self, - db: lancedb.DBConnection, all_records: list[dict], embedding_model: str, embedding_dim: int, @@ -1094,14 +1175,16 @@ def _add_chunk_records_unlocked( """Add the batch's chunk rows, writing meta on first use. Caller holds ``write_lock()``.""" if not all_records: return - ensure_table(db, CHUNKS_TABLE, self._chunks_schema()).add(all_records) + self._chunks_table().add(all_records) if self.get_meta() is None: self._write_meta_unlocked(embedding_model=embedding_model, embedding_dim=embedding_dim) def _batch_source_rows(self, items: list[ChunkWrite]) -> list[dict]: """One ``_sources`` row per batched document.""" return [ - self._source_row(it.source, it.file_hash, len(it.records), it.source_type, it.stat) + self._source_row( + it.source, it.file_hash, len(it.records), it.source_type, it.stat, it.meta + ) for it in items ] diff --git a/src/lilbee/data/store/fusion.py b/src/lilbee/data/store/fusion.py index 157784f47..1794fae3f 100644 --- a/src/lilbee/data/store/fusion.py +++ b/src/lilbee/data/store/fusion.py @@ -1,12 +1,15 @@ -"""Reciprocal-rank fusion of the vector and BM25 retrieval arms. +"""Reciprocal-rank fusion of the vector, BM25, and optional title arms. Each arm contributes ``(K + 1) / (K + rank)`` for the rows it retrieved (rank is 1-based, K = 60, the standard RRF constant); a row's canonical -score is the mean of its two arm contributions, so it lives in [0, 1] and -a row ranked first by both arms scores exactly 1. Rows seen by only one -arm score half of that arm's contribution, which still places an arm's -top hit above every row deep in the other arm: the property that keeps a -lexical-only identifier match visible next to dense neighbors. +score is the weight-normalized sum of its arm contributions, so it lives +in [0, 1] and a row ranked first by every arm scores exactly 1. Rows seen +by only one arm score that arm's share of the total weight, which still +places an arm's top hit above every row deep in the other arms: the +property that keeps a lexical-only identifier match visible next to dense +neighbors. The vector and chunk-BM25 arms weigh 1 each; the optional +title arm weighs ``title_weight``, so enabling it rescales the shares +without leaving the canonical range. Rank fusion is deliberate. A convex combination of normalized raw scores (``alpha * vector_similarity + (1 - alpha) * normalized_bm25``) was tried @@ -55,28 +58,49 @@ def _rank_weight(rank: int) -> float: return (_RRF_K + 1) / (_RRF_K + rank) +def _merge_arm( + merged: dict[tuple[str, int], SearchChunk], + rows: list[SearchChunk], + share: float, +) -> None: + """Fold one arm's ranked rows into *merged*, each contributing *share* of its rank weight. + + A lexical row (title or chunk FTS) carries ``bm25_score``; when the row was + already seen, that provenance is kept from whichever lexical arm set it + first, so the lexical-support exemption applies either way. + """ + for rank, row in enumerate(rows, start=1): + key = _key(row) + contribution = _rank_weight(rank) * share + seen = merged.get(key) + if seen is None: + merged[key] = row.model_copy(update={"score": contribution}) + else: + update: dict[str, object] = {"score": (seen.score or 0.0) + contribution} + if seen.bm25_score is None and row.bm25_score is not None: + update["bm25_score"] = row.bm25_score + merged[key] = seen.model_copy(update=update) + + def fuse_arms( vector_rows: list[SearchChunk], fts_rows: list[SearchChunk], + title_rows: list[SearchChunk] | None = None, + *, + title_weight: float = 1.0, ) -> list[SearchChunk]: - """Merge the two arms into one list scored by reciprocal rank. + """Merge the arms into one list scored by reciprocal rank. - Both arms weigh equally. Rows found by both arms carry both provenance - fields (``distance`` from the vector arm, ``bm25_score`` from the FTS - arm). The result is sorted by ``score`` descending and deduplicated on - ``(source, chunk_index)``. + The vector and chunk-FTS arms weigh equally; a non-empty *title_rows* arm + joins at *title_weight* relative to them. Rows found by several arms carry + every provenance field (``distance`` from the vector arm, ``bm25_score`` + from the FTS arms). The result is sorted by ``score`` descending and + deduplicated on ``(source, chunk_index)``. """ + total_weight = 2.0 + (title_weight if title_rows else 0.0) merged: dict[tuple[str, int], SearchChunk] = {} - for rank, row in enumerate(vector_rows, start=1): - merged[_key(row)] = row.model_copy(update={"score": _rank_weight(rank) / 2}) - for rank, row in enumerate(fts_rows, start=1): - key = _key(row) - lexical = _rank_weight(rank) / 2 - seen = merged.get(key) - if seen is None: - merged[key] = row.model_copy(update={"score": lexical}) - else: - merged[key] = seen.model_copy( - update={"score": (seen.score or 0.0) + lexical, "bm25_score": row.bm25_score} - ) + _merge_arm(merged, vector_rows, 1.0 / total_weight) + _merge_arm(merged, fts_rows, 1.0 / total_weight) + if title_rows: + _merge_arm(merged, title_rows, title_weight / total_weight) return sorted(merged.values(), key=lambda r: r.score or 0.0, reverse=True) diff --git a/src/lilbee/data/store/lance_helpers.py b/src/lilbee/data/store/lance_helpers.py index 6565f3ac1..a4b3a69a5 100644 --- a/src/lilbee/data/store/lance_helpers.py +++ b/src/lilbee/data/store/lance_helpers.py @@ -128,11 +128,11 @@ def _chunk_type_predicate(chunk_type: ChunkType | str) -> str: return f"chunk_type = '{escaped}'" -def _has_fts_index(table: lancedb.table.Table) -> bool: - """Return True when an FTS index on the chunk column already exists.""" +def _has_fts_index(table: lancedb.table.Table, column: str = "chunk") -> bool: + """Return True when an FTS index on *column* already exists.""" try: for idx in table.list_indices(): - if idx.index_type == "FTS" and "chunk" in idx.columns: + if idx.index_type == "FTS" and column in idx.columns: return True except Exception: return False diff --git a/src/lilbee/data/store/schema.py b/src/lilbee/data/store/schema.py index 21bd15324..60265911c 100644 --- a/src/lilbee/data/store/schema.py +++ b/src/lilbee/data/store/schema.py @@ -38,6 +38,9 @@ def _sources_schema() -> pa.Schema: pa.field("size_bytes", pa.int64()), pa.field("mtime_ns", pa.int64()), pa.field("stat_captured_ns", pa.int64()), + pa.field("title", pa.utf8()), + pa.field("authors", pa.utf8()), + pa.field("created_at", pa.utf8()), ] ) diff --git a/src/lilbee/data/store/types.py b/src/lilbee/data/store/types.py index f7f0cf51a..2d380825f 100644 --- a/src/lilbee/data/store/types.py +++ b/src/lilbee/data/store/types.py @@ -45,6 +45,19 @@ class SourceType(StrEnum): IMPORTED = "imported" +class SourceMeta(NamedTuple): + """Document-level metadata captured at extraction time. + + ``title`` is always derivable (extraction metadata or the cleaned filename + stem); ``authors`` and ``created_at`` are only present when the extractor + reports them. Empty strings persist as NULL so old and new rows read alike. + """ + + title: str = "" + authors: str = "" + created_at: str = "" + + class ChunkWrite(NamedTuple): """One document's chunks plus its source-table update, for a batched write. @@ -62,6 +75,7 @@ class ChunkWrite(NamedTuple): stat: SourceStat | None = None page_texts: list[dict] | None = None source_type: SourceType = SourceType.DOCUMENT + meta: SourceMeta | None = None class ChunkType(StrEnum): @@ -152,6 +166,10 @@ def _coerce_none_chunk_type(cls, v: str | None) -> str: """LanceDB rows from before the chunk_type column was added return None.""" return v if v is not None else ChunkType.RAW + # Document title at ingest time; None on rows written before the column + # existed (or by writers that carry no title, e.g. wiki pages). + title: str | None = None + page_start: int page_end: int line_start: int @@ -192,6 +210,11 @@ class SourceRecord(TypedDict): size_bytes: NotRequired[int] mtime_ns: NotRequired[int] stat_captured_ns: NotRequired[int] + # Extraction-time document metadata; absent or None on rows written + # before the columns existed, and None when the extractor reported none. + title: NotRequired[str | None] + authors: NotRequired[str | None] + created_at: NotRequired[str | None] # Sentinel for the stat columns on rows written before they existed (or for diff --git a/tests/server/test_handlers.py b/tests/server/test_handlers.py index a3c8d29f1..060b6120b 100644 --- a/tests/server/test_handlers.py +++ b/tests/server/test_handlers.py @@ -86,6 +86,7 @@ def _make_kreuzberg_result(text: str = "Some extracted text. " * 20, num_chunks: result = mock.MagicMock() result.chunks = chunks result.content = text + result.metadata = {} return result diff --git a/tests/test_export.py b/tests/test_export.py index 2b9259540..a022796cd 100644 --- a/tests/test_export.py +++ b/tests/test_export.py @@ -263,3 +263,12 @@ def counting_batch(items): assert legacy == [] # none of the old per-op unlocked writes were used # The atomic write still landed the source as IMPORTED. assert store.get_sources()[0]["source_type"] == SourceType.IMPORTED + + async def test_import_stamps_stem_title(self, services): + # Imports carry no extraction metadata; the stem-derived title keeps + # imported chunks visible to the title search arm. + store = services + await import_dataset(store, [_page("field_notes.pdf", 1, "page body")]) + chunks = store.get_chunks_by_source("field_notes.pdf") + assert chunks and all(c.title == "field notes" for c in chunks) + assert store.get_sources()[0]["title"] == "field notes" diff --git a/tests/test_formats.py b/tests/test_formats.py index d0fcd97f5..bbfef28df 100644 --- a/tests/test_formats.py +++ b/tests/test_formats.py @@ -68,6 +68,7 @@ def _make_kreuzberg_result(text="Extracted content. " * 10, num_chunks=1): result = mock.MagicMock() result.chunks = chunks result.content = text + result.metadata = {} return result diff --git a/tests/test_ingest.py b/tests/test_ingest.py index b605d8633..66043995e 100644 --- a/tests/test_ingest.py +++ b/tests/test_ingest.py @@ -161,6 +161,7 @@ def _make_kreuzberg_result( result.chunks = chunks result.content = text result.document = document + result.metadata = {} result.pages = ( [{"page_number": i + 1, "content": chunks[i].content} for i in range(num_chunks)] if has_pages @@ -175,6 +176,7 @@ def _make_empty_result(): result.chunks = [] result.content = "" result.document = None + result.metadata = {} return result @@ -1090,7 +1092,14 @@ class TestZeroChunkPageTextPersistence: @staticmethod async def _pages_no_chunks( - path, source_name, content_type, *, quiet=False, on_progress=None, page_texts_out=None + path, + source_name, + content_type, + *, + quiet=False, + on_progress=None, + page_texts_out=None, + meta_out=None, ): if page_texts_out is not None: page_texts_out.append( @@ -3184,3 +3193,52 @@ def test_no_marker_when_nothing_removed(self, isolated_env, mock_svc): mock_svc.store.remove_documents.return_value = RemoveResult(removed=[], not_found=["gone"]) remove_documents_durably(["gone"]) assert load_skip_markers(cfg.data_root) == {} + + +class TestTitleStamping: + """Every produced record carries the document title; the source row its metadata.""" + + @mock.patch("kreuzberg.extract_file_sync", new_callable=Mock) + async def test_extracted_metadata_flows_to_records_and_source_row( + self, mock_kf, isolated_env, mock_svc + ): + result = _make_kreuzberg_result() + result.metadata = { + "title": "Extracted Title", + "authors": ["Ada", "Grace"], + "created_at": "2020-01-01", + } + mock_kf.return_value = result + (isolated_env / "report_2021.txt").write_text("body text") + from lilbee.data.ingest import sync + + await sync(quiet=True) + items = mock_svc.store.write_chunks_batch.call_args.args[0] + item = next(it for it in items if it.source == "report_2021.txt") + assert item.meta.title == "Extracted Title" + assert item.meta.authors == "Ada, Grace" + assert item.meta.created_at == "2020-01-01" + assert all(r["title"] == "Extracted Title" for r in item.records) + + @mock.patch("kreuzberg.extract_file_sync", new_callable=Mock) + async def test_missing_metadata_falls_back_to_stem(self, mock_kf, isolated_env, mock_svc): + mock_kf.return_value = _make_kreuzberg_result() + (isolated_env / "annual_wildlife_survey.txt").write_text("body text") + from lilbee.data.ingest import sync + + await sync(quiet=True) + items = mock_svc.store.write_chunks_batch.call_args.args[0] + item = next(it for it in items if it.source == "annual_wildlife_survey.txt") + assert item.meta.title == "annual wildlife survey" + assert item.meta.authors == "" + assert all(r["title"] == "annual wildlife survey" for r in item.records) + + async def test_markdown_title_derives_from_stem(self, isolated_env, mock_svc): + (isolated_env / "meeting_notes.md").write_text("# Heading\n\nSome content here.") + from lilbee.data.ingest import sync + + await sync(quiet=True) + items = mock_svc.store.write_chunks_batch.call_args.args[0] + item = next(it for it in items if it.source == "meeting_notes.md") + assert item.meta.title == "meeting notes" + assert all(r["title"] == "meeting notes" for r in item.records) diff --git a/tests/test_ingest_title.py b/tests/test_ingest_title.py new file mode 100644 index 000000000..a3ebccbf8 --- /dev/null +++ b/tests/test_ingest_title.py @@ -0,0 +1,57 @@ +"""Tests for document title and source-metadata derivation at ingest.""" + +from lilbee.data.ingest.title import derive_title, source_meta_from_extraction + + +class TestDeriveTitle: + def test_metadata_title_wins(self): + assert derive_title("q3_report.pdf", "Quarterly Revenue Report") == ( + "Quarterly Revenue Report" + ) + + def test_metadata_title_is_stripped(self): + assert derive_title("a.pdf", " Padded Title \n") == "Padded Title" + + def test_blank_metadata_falls_back_to_stem(self): + assert derive_title("survey_214.pdf", " ") == "survey 214" + + def test_none_metadata_falls_back_to_stem(self): + assert derive_title("annual-wildlife-survey.pdf") == "annual wildlife survey" + + def test_stem_keeps_meaningful_dots(self): + assert derive_title("spec-v3.0.pdf") == "spec v3.0" + + def test_subdirectory_source_uses_basename_stem(self): + assert derive_title("filings/2024/notice_of_appeal.pdf") == "notice of appeal" + + def test_mixed_separators_collapse(self): + assert derive_title("a_b-c d.txt") == "a b c d" + + +class TestSourceMetaFromExtraction: + def test_full_metadata(self): + meta = source_meta_from_extraction( + {"title": "The Title", "authors": ["Ada", "Grace"], "created_at": "2021-05-01"}, + "x.pdf", + ) + assert meta.title == "The Title" + assert meta.authors == "Ada, Grace" + assert meta.created_at == "2021-05-01" + + def test_empty_metadata_derives_stem_title(self): + meta = source_meta_from_extraction({}, "field_notes.pdf") + assert meta.title == "field notes" + assert meta.authors == "" + assert meta.created_at == "" + + def test_falsy_authors_are_dropped(self): + meta = source_meta_from_extraction({"authors": ["Ada", "", None]}, "x.pdf") + assert meta.authors == "Ada" + + def test_none_values_tolerated(self): + meta = source_meta_from_extraction( + {"title": None, "authors": None, "created_at": None}, "notes.md" + ) + assert meta.title == "notes" + assert meta.authors == "" + assert meta.created_at == "" diff --git a/tests/test_settings.py b/tests/test_settings.py index 30f915e51..ec5541174 100644 --- a/tests/test_settings.py +++ b/tests/test_settings.py @@ -478,3 +478,31 @@ def test_win32_platform_sim_does_not_break_save(self, tmp_path, monkeypatch) -> settings.save(tmp_path, {"key": "value", "unicode": "é"}) result = settings.load(tmp_path) assert result["unicode"] == "é" + + +class TestTitleSearchSettings: + """The title-arm knobs are exposed on every settings surface.""" + + def test_title_search_in_settings_map(self): + from lilbee.app.settings_map import SETTINGS_MAP, get_default + + defn = SETTINGS_MAP["title_search"] + assert defn.writable is True + assert defn.type is bool + assert defn.group == "Retrieval" + assert get_default("title_search") is False + + def test_title_search_weight_in_settings_map(self): + from lilbee.app.settings_map import SETTINGS_MAP, get_default + + defn = SETTINGS_MAP["title_search_weight"] + assert defn.writable is True + assert defn.type is float + assert defn.group == "Retrieval" + assert get_default("title_search_weight") == 0.5 + + def test_title_search_fields_are_writable_for_programmatic_surfaces(self): + from lilbee.config_meta import WRITABLE_CONFIG_FIELDS + + assert "title_search" in WRITABLE_CONFIG_FIELDS + assert "title_search_weight" in WRITABLE_CONFIG_FIELDS diff --git a/tests/test_store.py b/tests/test_store.py index 73514aa36..3806b12bd 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -163,7 +163,7 @@ def test_second_call_optimizes_instead_of_rebuilding(self, store): optimize_spy.assert_called_once() def test_first_call_creates_without_replace(self, store): - """Fresh table goes through create_fts_index path with replace=False.""" + """Fresh table creates the chunk and title indexes, both with replace=False.""" store.add_chunks(_make_records()) table = store.open_table("chunks") assert table is not None @@ -171,10 +171,9 @@ def test_first_call_creates_without_replace(self, store): with mock.patch.object(type(table), "create_fts_index") as create_spy: store.ensure_fts_index() - create_spy.assert_called_once() + assert [c.args[0] for c in create_spy.call_args_list] == ["chunk", "title"] # Verify replace was NOT True (would defeat the purpose of incremental) - _args, kwargs = create_spy.call_args - assert kwargs.get("replace") is False + assert all(c.kwargs.get("replace") is False for c in create_spy.call_args_list) def test_bm25_probe_populates_bm25_score(self, store): """LanceDB FTS returns rows keyed on ``_score``; the probe must surface it as @@ -197,6 +196,15 @@ def test_bm25_probe_filters_by_chunk_type(self, store): assert results assert all(r.chunk_type == ChunkType.WIKI for r in results) + def test_bm25_probe_survives_a_failing_search(self, store): + """A probe whose LanceDB query raises degrades to no hits, not an error.""" + store.add_chunks(_make_records()) + store.ensure_fts_index() + table = store.open_table("chunks") + assert table is not None + with mock.patch.object(type(table), "search", side_effect=RuntimeError("boom")): + assert store.bm25_probe("text") == [] + class TestSearchChunkScoreAlias: @staticmethod @@ -2195,3 +2203,205 @@ def test_negative_row_count_clamps_to_floor(self): from lilbee.data.store.core import _ANN_NPROBES_FLOOR, _ann_nprobes assert _ann_nprobes(-5) == _ANN_NPROBES_FLOOR + + +def _titled_records(source, n, *, title, base=0.1, dim=None): + """Chunk records for one source with a document title on every row.""" + if dim is None: + dim = cfg.embedding_dim + return [ + { + "source": source, + "content_type": "text", + "chunk_type": "raw", + "page_start": 0, + "page_end": 0, + "line_start": 0, + "line_end": 0, + "chunk": f"plain body words {source} {i}", + "chunk_index": i, + "title": title, + "vector": [base + i / 100] * dim, + } + for i in range(n) + ] + + +def _create_pre_title_chunks_table(store): + """Create the chunks table with the pre-title schema, as an old index would have.""" + import pyarrow as pa + + schema = pa.schema( + [ + pa.field("source", pa.utf8()), + pa.field("content_type", pa.utf8()), + pa.field("chunk_type", pa.utf8()), + pa.field("page_start", pa.int32()), + pa.field("page_end", pa.int32()), + pa.field("line_start", pa.int32()), + pa.field("line_end", pa.int32()), + pa.field("chunk", pa.utf8()), + pa.field("chunk_index", pa.int32()), + pa.field("vector", pa.list_(pa.float32(), cfg.embedding_dim)), + ] + ) + return store.get_db().create_table("chunks", schema=schema) + + +class TestTitleSearch: + """The title lexical arm: BM25 over document titles fused into hybrid search.""" + + def test_title_arm_surfaces_title_only_match(self, store, test_config): + """A term that lives only in a document's title reaches hybrid results + with lexical support (bm25_score) when title_search is on.""" + test_config.title_search = True + store.add_chunks(_titled_records("a.pdf", 2, title="zebra manifesto", base=0.9)) + store.add_chunks(_titled_records("b.pdf", 2, title="meeting notes", base=0.1)) + store.ensure_fts_index() + query_vec = [0.1] * test_config.embedding_dim + results = store.search(query_vec, top_k=4, max_distance=0, query_text="zebra") + matched = [r for r in results if r.source == "a.pdf"] + assert matched + assert all(r.bm25_score is not None for r in matched) + + def test_title_arm_off_by_default(self, store, test_config): + """With title_search off, a title-only term earns no lexical support.""" + assert test_config.title_search is False + store.add_chunks(_titled_records("a.pdf", 2, title="zebra manifesto")) + store.ensure_fts_index() + query_vec = [0.1] * test_config.embedding_dim + results = store.search(query_vec, top_k=4, max_distance=0, query_text="zebra") + assert all(r.bm25_score is None for r in results) + + def test_title_arm_respects_chunk_type_filter(self, store, test_config): + from lilbee.data.store.lance_helpers import _has_fts_index + + store.add_chunks(_titled_records("a.pdf", 2, title="zebra manifesto")) + store.ensure_fts_index() + table = store.open_table("chunks") + assert _has_fts_index(table, "title") + assert store._title_arm(table, "zebra", 5, ChunkType.RAW) + assert store._title_arm(table, "zebra", 5, ChunkType.WIKI) == [] + + def test_old_index_without_title_column_still_searches(self, store, test_config): + """Feature detection: a pre-title index searches fine and the title arm + silently contributes nothing.""" + test_config.title_search = True + table = _create_pre_title_chunks_table(store) + table.add(_make_records()) + store.ensure_fts_index() + assert store._title_arm(table, "chunk", 5, None) == [] + query_vec = [0.5] * test_config.embedding_dim + results = store.search(query_vec, top_k=3, max_distance=0, query_text="chunk number") + assert results + assert all(r.title is None for r in results) + + def test_add_chunks_evolves_pre_title_table(self, store): + """A write to an old index adds the title column in place.""" + table = _create_pre_title_chunks_table(store) + table.add(_make_records()) + assert "title" not in table.schema.names + store.add_chunks(_titled_records("new.pdf", 1, title="fresh document")) + assert "title" in store.open_table("chunks").schema.names + rows = store.get_chunks_by_source("new.pdf") + assert [r.title for r in rows] == ["fresh document"] + + def test_bm25_probe_stays_chunk_scoped(self, store): + """The probe pins the chunk column: a title-only term is not a probe hit.""" + store.add_chunks(_titled_records("a.pdf", 2, title="zebra manifesto")) + store.ensure_fts_index() + assert store.bm25_probe("zebra") == [] + assert store.bm25_probe("plain body words") + + def test_title_index_failure_never_blocks_chunk_index(self, store, test_config): + """A failing title index leaves chunk FTS ready; the arm degrades to empty.""" + test_config.title_search = True + store.add_chunks(_titled_records("a.pdf", 1, title="zebra manifesto")) + table = store.open_table("chunks") + real_create = type(table).create_fts_index + + def _fail_title(self, column, **kwargs): + if column == "title": + raise RuntimeError("boom") + return real_create(self, column, **kwargs) + + with mock.patch.object(type(table), "create_fts_index", _fail_title): + store.ensure_fts_index() + assert store._fts_ready + assert store._title_arm(store.open_table("chunks"), "zebra", 5, None) == [] + + +class TestSourceMetadata: + """Extraction-time document metadata persisted on the sources table.""" + + def test_upsert_source_persists_meta(self, store): + from lilbee.data.store import SourceMeta + + store.upsert_source( + "a.pdf", + "hash1", + 3, + meta=SourceMeta(title="The Title", authors="Ada, Grace", created_at="2021-05-01"), + ) + row = store.get_sources()[0] + assert row["title"] == "The Title" + assert row["authors"] == "Ada, Grace" + assert row["created_at"] == "2021-05-01" + + def test_absent_meta_persists_null(self, store): + store.upsert_source("a.pdf", "hash1", 3) + row = store.get_sources()[0] + assert row["title"] is None + assert row["authors"] is None + assert row["created_at"] is None + + def test_pre_meta_sources_table_evolves_in_place(self, store): + """An old sources table gains the metadata columns on the next write.""" + import pyarrow as pa + + from lilbee.data.store import SourceMeta + + old_schema = pa.schema( + [ + pa.field("filename", pa.utf8()), + pa.field("file_hash", pa.utf8()), + pa.field("ingested_at", pa.utf8()), + pa.field("chunk_count", pa.int32()), + pa.field("source_type", pa.utf8()), + ] + ) + table = store.get_db().create_table("_sources", schema=old_schema) + table.add( + [ + { + "filename": "old.pdf", + "file_hash": "h0", + "ingested_at": "2020-01-01T00:00:00+00:00", + "chunk_count": 1, + "source_type": "document", + } + ] + ) + store.upsert_source("new.pdf", "h1", 2, meta=SourceMeta(title="New Doc")) + rows = {r["filename"]: r for r in store.get_sources()} + assert rows["old.pdf"]["title"] is None + assert rows["new.pdf"]["title"] == "New Doc" + + def test_batched_write_persists_meta(self, store): + from lilbee.data.store import ChunkWrite, SourceMeta + + records = _titled_records("doc.pdf", 1, title="Batched Title") + store.write_chunks_batch( + [ + ChunkWrite( + "doc.pdf", + "h", + records, + needs_cleanup=False, + meta=SourceMeta(title="Batched Title", authors="Ada"), + ) + ] + ) + row = store.get_sources()[0] + assert row["title"] == "Batched Title" + assert row["authors"] == "Ada" diff --git a/tests/test_store_fusion.py b/tests/test_store_fusion.py index 5e20e68e5..82d44009b 100644 --- a/tests/test_store_fusion.py +++ b/tests/test_store_fusion.py @@ -186,3 +186,59 @@ def test_drops_vector_only_far_row_keeps_supported(self): far_supported = _chunk("identifier.md", 0, distance=1.4, bm25=25.0) kept = _drop_unsupported_far_rows([far_unsupported, far_supported], 0.75) assert [r.source for r in kept] == ["identifier.md"] + + +class TestTitleArmFusion: + """The optional third arm: BM25 over document titles, weight-normalized.""" + + def test_top_of_all_three_arms_scores_one(self): + fused = fuse_arms( + [_chunk("a.md", 0, distance=0.3)], + [_chunk("a.md", 0, bm25=12.0)], + [_chunk("a.md", 0, bm25=3.0)], + title_weight=0.5, + ) + assert fused[0].score == pytest.approx(1.0) + + def test_title_only_row_scores_its_weight_share(self): + weight = 0.5 + fused = fuse_arms([], [], [_chunk("t.md", 0, bm25=3.0)], title_weight=weight) + assert fused[0].score == pytest.approx(weight / (2.0 + weight)) + + def test_empty_title_arm_matches_two_arm_scores(self): + """No title rows = the classic two-arm fusion, share-for-share.""" + vector = [_chunk("a.md", 0, distance=0.3), _chunk("b.md", 0, distance=0.4)] + fts = [_chunk("a.md", 0, bm25=12.0)] + two_arm = fuse_arms(vector, fts) + with_empty_title = fuse_arms(vector, fts, [], title_weight=0.5) + assert [(r.source, r.score) for r in two_arm] == [ + (r.source, r.score) for r in with_empty_title + ] + + def test_title_match_counts_as_lexical_support(self): + """A row only the title arm matched carries bm25_score, so the + distance exemption sees lexical support.""" + fused = fuse_arms( + [_chunk("a.md", 0, distance=1.5)], + [], + [_chunk("a.md", 0, bm25=4.0)], + title_weight=0.5, + ) + assert fused[0].bm25_score == pytest.approx(4.0) + assert fused[0].distance == pytest.approx(1.5) + + def test_chunk_arm_bm25_provenance_wins_over_title(self): + """When both lexical arms match a row, the first-seen bm25_score is kept.""" + fused = fuse_arms( + [], + [_chunk("a.md", 0, bm25=12.0)], + [_chunk("a.md", 0, bm25=3.0)], + title_weight=0.5, + ) + assert fused[0].bm25_score == pytest.approx(12.0) + + def test_title_weight_scales_contribution(self): + low = fuse_arms([], [], [_chunk("t.md", 0, bm25=3.0)], title_weight=0.2) + high = fuse_arms([], [], [_chunk("t.md", 0, bm25=3.0)], title_weight=1.0) + assert low[0].score < high[0].score + assert high[0].score == pytest.approx(1.0 / 3.0) From 02e4078d40c176941d0b0a9cbb8043221a33b7f7 Mon Sep 17 00:00:00 2001 From: Tobias Perelstein <5562156+tobocop2@users.noreply.github.com> Date: Fri, 17 Jul 2026 19:10:10 -0400 Subject: [PATCH 03/28] Coerce a string data_root into a Path before deriving child dirs Setting the data root through the LILBEE_DATA_ROOT environment variable delivered it to the config resolver as a raw string, and deriving the documents/data/lancedb child directories with the / operator then raised TypeError: unsupported operand type(s) for /: 'str' and 'str'. Coerce the value to Path up front so the env var works the same as every other way of setting the data root. --- src/lilbee/core/config/model.py | 5 ++++- tests/test_config.py | 11 +++++++++++ 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/src/lilbee/core/config/model.py b/src/lilbee/core/config/model.py index 6f4b0043d..aaee41ba7 100644 --- a/src/lilbee/core/config/model.py +++ b/src/lilbee/core/config/model.py @@ -966,7 +966,10 @@ def _resolve_defaults(cls, data: Any) -> Any: else: local = find_local_root() data["data_root"] = local if local is not None else default_data_dir() - root = data["data_root"] + # data_root may arrive as a raw string (e.g. from LILBEE_DATA_ROOT); the + # child-path derivations below use ``/``, so coerce to Path first. + root = Path(data["data_root"]) + data["data_root"] = root if data.get("documents_dir") in (None, _UNSET_PATH): data["documents_dir"] = root / "documents" if data.get("data_dir") in (None, _UNSET_PATH): diff --git a/tests/test_config.py b/tests/test_config.py index 7c391100c..7158e4241 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -747,6 +747,17 @@ def test_env_overrides_toml(self, tmp_path) -> None: with mock.patch.dict(os.environ, env, clear=True): assert Config().semantic_chunking is True + def test_data_root_env_var_is_coerced_to_path(self, tmp_path) -> None: + """LILBEE_DATA_ROOT sets the data_root field directly as a string; + deriving the child paths from it must not raise on str / str.""" + env = _clean_env() + env["LILBEE_DATA_ROOT"] = str(tmp_path) + with mock.patch.dict(os.environ, env, clear=True): + c = Config() + assert c.data_root == tmp_path + assert c.documents_dir == tmp_path / "documents" + assert c.data_dir == tmp_path / "data" + class TestTopicThresholdConfig: def test_default_is_0_75(self, tmp_path) -> None: From ded0440620256ea1a6e28cf3e294d6a0ad190ad2 Mon Sep 17 00:00:00 2001 From: Tobias Perelstein <5562156+tobocop2@users.noreply.github.com> Date: Fri, 17 Jul 2026 17:18:38 -0400 Subject: [PATCH 04/28] Build the FTS index with token positions so phrase queries work lilbee passes raw user query text to the lexical arm, so a query LanceDB parses as a phrase (e.g. a quoted span) reaches it as a phrase query. The chunk and title FTS indexes were created without positions, so those queries raised and the arm silently returned nothing. Create both indexes with_position=True. Existing indexes need a reindex to gain phrase support; they degrade to vector-only until then rather than erroring. bb-e4qw (cherry picked from commit a2cd68432e04978d577a7880de8ee31f8f214c7e) --- src/lilbee/data/store/core.py | 7 +++++-- tests/test_store.py | 16 ++++++++++++++++ 2 files changed, 21 insertions(+), 2 deletions(-) diff --git a/src/lilbee/data/store/core.py b/src/lilbee/data/store/core.py index 56e1bcefd..e3ad8608a 100644 --- a/src/lilbee/data/store/core.py +++ b/src/lilbee/data/store/core.py @@ -432,7 +432,10 @@ def ensure_fts_index(self) -> None: table.optimize() log.debug("FTS index optimized on '%s'", CHUNKS_TABLE) else: - table.create_fts_index("chunk", replace=False) + # with_position lets the lexical arm serve phrase queries; + # raw user query text can reach LanceDB as a phrase, which + # errors on a positionless index. + table.create_fts_index("chunk", replace=False, with_position=True) log.debug("FTS index created on '%s'", CHUNKS_TABLE) self._ensure_title_fts_unlocked(table) self._fts_ready = True @@ -448,7 +451,7 @@ def _ensure_title_fts_unlocked(self, table: lancedb.table.Table) -> None: if _TITLE_COLUMN not in table.schema.names or _has_fts_index(table, _TITLE_COLUMN): return try: - table.create_fts_index(_TITLE_COLUMN, replace=False) + table.create_fts_index(_TITLE_COLUMN, replace=False, with_position=True) log.debug("Title FTS index created on '%s'", CHUNKS_TABLE) except Exception: log.debug("Title FTS index create failed", exc_info=True) diff --git a/tests/test_store.py b/tests/test_store.py index 3806b12bd..4a67c41d0 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -174,6 +174,22 @@ def test_first_call_creates_without_replace(self, store): assert [c.args[0] for c in create_spy.call_args_list] == ["chunk", "title"] # Verify replace was NOT True (would defeat the purpose of incremental) assert all(c.kwargs.get("replace") is False for c in create_spy.call_args_list) + # Both indexes carry token positions so the lexical arm can serve + # phrase queries; without this LanceDB raises on any phrase. + assert all(c.kwargs.get("with_position") is True for c in create_spy.call_args_list) + + def test_fts_phrase_query_does_not_fail(self, store): + """A multi-word phrase query must return matches, not raise on a missing index. + + lilbee passes raw user query text to the lexical arm, so a quoted phrase + reaches LanceDB as a phrase query. Without positional indexing that raises + (and bm25_probe swallows it, returning nothing); with positions it matches. + """ + store.add_chunks(_make_records()) + store.ensure_fts_index() + results = store.bm25_probe('"some text"') + assert results + assert all("some text" in r.chunk for r in results) def test_bm25_probe_populates_bm25_score(self, store): """LanceDB FTS returns rows keyed on ``_score``; the probe must surface it as From cac0a63ccea89be1c6520ce9968717ad20f19af4 Mon Sep 17 00:00:00 2001 From: Tobias Perelstein <5562156+tobocop2@users.noreply.github.com> Date: Fri, 17 Jul 2026 19:12:03 -0400 Subject: [PATCH 05/28] Make the BM25 fusion arm weight configurable Rank fusion has weighed the vector and BM25 arms equally, which dilutes a strong dense embedder on corpora where the lexical arm adds more noise than signal. Add lexical_fusion_weight (default 1.0, so the equal-weight behaviour is byte-for-byte unchanged) to scale the BM25 arm's share of the fused score down toward the dense arm. The right value is corpus-dependent and meant to be set from the retrieval benchmark, not guessed. Wired through the config field, its HTTP/MCP/env surfaces, the TUI settings map, and fuse_arms. --- docs/usage.md | 1 + src/lilbee/app/settings_map.py | 6 ++++++ src/lilbee/core/config/model.py | 7 +++++++ src/lilbee/data/store/core.py | 1 + src/lilbee/data/store/fusion.py | 23 +++++++++++++---------- tests/test_store_fusion.py | 22 ++++++++++++++++++++++ 6 files changed, 50 insertions(+), 10 deletions(-) diff --git a/docs/usage.md b/docs/usage.md index 8b2d2a2f6..7d2917ca1 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -750,6 +750,7 @@ something feels off. | `LILBEE_HYDE_WEIGHT` | `0.7` | How much to trust HyDE results relative to the direct query (0.0-1.0) | | `LILBEE_TITLE_SEARCH` | `false` | Match queries against document titles as a third hybrid-search arm, so a query naming a document by title surfaces its chunks | | `LILBEE_TITLE_SEARCH_WEIGHT` | `0.5` | Title arm weight in rank fusion (1.0 = equal voice with the vector and text arms) | +| `LILBEE_LEXICAL_FUSION_WEIGHT` | `1.0` | BM25 arm weight in rank fusion (1.0 = equal to the vector arm; lower it when a strong embedder should dominate a corpus where the lexical arm adds noise) | | `LILBEE_ADAPTIVE_THRESHOLD` | `false` | Widen the distance threshold step by step when too few results pass, on the vector-only fallback path (no effect on the default hybrid search) | | `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 | diff --git a/src/lilbee/app/settings_map.py b/src/lilbee/app/settings_map.py index 43f1c4ee8..95fa5d58f 100644 --- a/src/lilbee/app/settings_map.py +++ b/src/lilbee/app/settings_map.py @@ -829,6 +829,12 @@ def get_default(key: str) -> object: group=SettingGroup.RETRIEVAL, help_text="Title arm weight in rank fusion (1.0 = equal voice with the other arms)", ), + "lexical_fusion_weight": SettingDef( + float, + nullable=False, + group=SettingGroup.RETRIEVAL, + help_text="BM25 arm weight in fusion (1.0 = equal to vector; lower to favor dense)", + ), "history_rewrite": SettingDef( bool, nullable=False, diff --git a/src/lilbee/core/config/model.py b/src/lilbee/core/config/model.py index 642473975..0ad011a9f 100644 --- a/src/lilbee/core/config/model.py +++ b/src/lilbee/core/config/model.py @@ -203,6 +203,13 @@ class Config(BaseSettings): # with the vector and chunk arms). title_search_weight: float = ConfigField(default=0.5, ge=0.0, le=1.0, writable=True) + # Lexical (BM25) arm weight relative to the vector arm in rank fusion. + # 1.0 keeps the two arms equal (the historical behaviour); lowering it lets + # a strong dense embedder dominate on corpora where the lexical arm adds + # noise rather than signal. The right value is corpus-dependent and set by + # the retrieval benchmark, not guessed here. + lexical_fusion_weight: float = ConfigField(default=1.0, ge=0.0, le=1.0, writable=True) + # Chunk count at/above which sync builds an approximate (ANN) vector index # so search stays fast at millions of vectors. Below this, search uses exact # flat scan (faster and exact for small vaults). 0 disables the ANN index. diff --git a/src/lilbee/data/store/core.py b/src/lilbee/data/store/core.py index e3ad8608a..0e94997b0 100644 --- a/src/lilbee/data/store/core.py +++ b/src/lilbee/data/store/core.py @@ -704,6 +704,7 @@ def _hybrid_search( self._vector_arm(table, query_vector, top_k, chunk_type), self._fts_arm(table, query_text, top_k, chunk_type), title_rows, + lexical_weight=self._config.lexical_fusion_weight, title_weight=self._config.title_search_weight, ) fused = _drop_unsupported_far_rows(fused, max_distance) diff --git a/src/lilbee/data/store/fusion.py b/src/lilbee/data/store/fusion.py index 1794fae3f..8804a4a3f 100644 --- a/src/lilbee/data/store/fusion.py +++ b/src/lilbee/data/store/fusion.py @@ -7,9 +7,10 @@ by only one arm score that arm's share of the total weight, which still places an arm's top hit above every row deep in the other arms: the property that keeps a lexical-only identifier match visible next to dense -neighbors. The vector and chunk-BM25 arms weigh 1 each; the optional -title arm weighs ``title_weight``, so enabling it rescales the shares -without leaving the canonical range. +neighbors. The vector arm weighs 1; the chunk-BM25 arm weighs +``lexical_weight`` (1.0 = equal voice, lower lets a strong dense arm +dominate); the optional title arm weighs ``title_weight``. Weights rescale +the shares without leaving the canonical range. Rank fusion is deliberate. A convex combination of normalized raw scores (``alpha * vector_similarity + (1 - alpha) * normalized_bm25``) was tried @@ -87,20 +88,22 @@ def fuse_arms( fts_rows: list[SearchChunk], title_rows: list[SearchChunk] | None = None, *, + lexical_weight: float = 1.0, title_weight: float = 1.0, ) -> list[SearchChunk]: """Merge the arms into one list scored by reciprocal rank. - The vector and chunk-FTS arms weigh equally; a non-empty *title_rows* arm - joins at *title_weight* relative to them. Rows found by several arms carry - every provenance field (``distance`` from the vector arm, ``bm25_score`` - from the FTS arms). The result is sorted by ``score`` descending and - deduplicated on ``(source, chunk_index)``. + The vector arm weighs 1; the chunk-FTS (lexical) arm weighs *lexical_weight* + relative to it (1.0 = equal voice, lower lets a strong dense arm dominate); + a non-empty *title_rows* arm joins at *title_weight*. Rows found by several + arms carry every provenance field (``distance`` from the vector arm, + ``bm25_score`` from the FTS arms). The result is sorted by ``score`` + descending and deduplicated on ``(source, chunk_index)``. """ - total_weight = 2.0 + (title_weight if title_rows else 0.0) + total_weight = 1.0 + lexical_weight + (title_weight if title_rows else 0.0) merged: dict[tuple[str, int], SearchChunk] = {} _merge_arm(merged, vector_rows, 1.0 / total_weight) - _merge_arm(merged, fts_rows, 1.0 / total_weight) + _merge_arm(merged, fts_rows, lexical_weight / total_weight) if title_rows: _merge_arm(merged, title_rows, title_weight / total_weight) return sorted(merged.values(), key=lambda r: r.score or 0.0, reverse=True) diff --git a/tests/test_store_fusion.py b/tests/test_store_fusion.py index 82d44009b..ac3c352dc 100644 --- a/tests/test_store_fusion.py +++ b/tests/test_store_fusion.py @@ -67,6 +67,28 @@ def test_lexical_only_row_survives_with_score(self): lexical_row = next(r for r in fused if r.source == "catalog_482.pdf") assert lexical_row.score == pytest.approx(0.5) + +class TestLexicalFusionWeight: + """The BM25 arm's fusion weight can be lowered so a strong dense arm dominates.""" + + def _lex_row(self, weight: float): + vec = [_chunk("noise.md", 0, distance=0.5)] + lex = [_chunk("cat.pdf", 0, bm25=35.0)] + return next(r for r in fuse_arms(vec, lex, lexical_weight=weight) if r.source == "cat.pdf") + + def test_default_weight_is_the_historical_equal_voice(self): + vec = [_chunk("noise.md", 0, distance=0.5)] + lex = [_chunk("cat.pdf", 0, bm25=35.0)] + default = {r.source: r.score for r in fuse_arms(vec, lex)} + explicit = {r.source: r.score for r in fuse_arms(vec, lex, lexical_weight=1.0)} + assert default == explicit + + def test_zero_weight_silences_the_lexical_arm(self): + assert self._lex_row(0.0).score == pytest.approx(0.0) + + def test_lower_weight_shrinks_the_lexical_contribution(self): + assert self._lex_row(0.5).score < self._lex_row(1.0).score + def test_top_lexical_hit_outranks_all_but_the_top_dense_neighbor(self): """The pinpoint-document failure mode: an FTS-arm top hit unseen by the vector arm must rank above every vector row except at most the From 1bf62ceef0119af5060274889aa33570462557ec Mon Sep 17 00:00:00 2001 From: Tobias Perelstein <5562156+tobocop2@users.noreply.github.com> Date: Fri, 17 Jul 2026 21:59:04 -0400 Subject: [PATCH 06/28] Keep hybrid search alive when FTS optimize() fails on a large corpus ensure_fts_index() runs table.optimize() whenever the FTS index already exists. On a large corpus that call can hit a LanceDB encoding bug and raise (Arrow list offset overflow), and the failure fell through to leave _fts_ready False -- silently dropping every subsequent query to vector-only with no error surfaced. The index itself is complete and still serves queries, so mark hybrid ready before the best-effort optimize and downgrade an optimize() failure to a warning. Found while benchmarking retrieval: every hybrid config collapsed to the vector-only baseline until this was fixed. (cherry picked from commit 6aedffb4c857ee1e5d2a62bbf5f24e0762514c53) --- src/lilbee/data/store/core.py | 22 +++++++++++++++++----- tests/test_store.py | 18 ++++++++++++++++++ 2 files changed, 35 insertions(+), 5 deletions(-) diff --git a/src/lilbee/data/store/core.py b/src/lilbee/data/store/core.py index 0e94997b0..fa09b03f9 100644 --- a/src/lilbee/data/store/core.py +++ b/src/lilbee/data/store/core.py @@ -427,18 +427,30 @@ def ensure_fts_index(self) -> None: return try: if _has_fts_index(table): - # One optimize folds new rows into every index on the - # table, the title index included. - table.optimize() - log.debug("FTS index optimized on '%s'", CHUNKS_TABLE) + # The index exists and already serves queries, so mark hybrid + # ready BEFORE the optimize: optimize() only folds new rows and + # compacts, and on a large corpus it can hit a LanceDB encoding + # bug and raise. Letting that failure fall through would leave + # _fts_ready False and silently drop every query to vector-only. + self._fts_ready = True + try: + # One optimize folds new rows into every index on the + # table, the title index included. + table.optimize() + log.debug("FTS index optimized on '%s'", CHUNKS_TABLE) + except Exception: + log.warning( + "FTS optimize() failed; the existing index still serves hybrid search", + exc_info=True, + ) else: # with_position lets the lexical arm serve phrase queries; # raw user query text can reach LanceDB as a phrase, which # errors on a positionless index. table.create_fts_index("chunk", replace=False, with_position=True) + self._fts_ready = True log.debug("FTS index created on '%s'", CHUNKS_TABLE) self._ensure_title_fts_unlocked(table) - self._fts_ready = True except Exception: log.debug("FTS index ensure failed (empty table?)", exc_info=True) diff --git a/tests/test_store.py b/tests/test_store.py index 4a67c41d0..d7b877ec9 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -162,6 +162,24 @@ def test_second_call_optimizes_instead_of_rebuilding(self, store): create_spy.assert_not_called() optimize_spy.assert_called_once() + def test_optimize_failure_keeps_hybrid_ready(self, store): + """An optimize() crash on an already-built index (a LanceDB encoding + bug bites large corpora) must not disable hybrid search: the index + still serves queries, so _fts_ready stays True instead of silently + dropping every query to the vector-only fallback.""" + store.add_chunks(_make_records()) + store.ensure_fts_index() # builds the index + store._fts_ready = False # a fresh process is unaware the index exists yet + table = store.open_table("chunks") + assert table is not None + with mock.patch.object( + type(table), + "optimize", + side_effect=RuntimeError("lance list offset overflow"), + ): + store.ensure_fts_index() + assert store._fts_ready is True + def test_first_call_creates_without_replace(self, store): """Fresh table creates the chunk and title indexes, both with replace=False.""" store.add_chunks(_make_records()) From 0ea29eba93d05fa06addd65b0125ed6a3af0265f Mon Sep 17 00:00:00 2001 From: Tobias Perelstein <5562156+tobocop2@users.noreply.github.com> Date: Fri, 17 Jul 2026 23:18:51 -0400 Subject: [PATCH 07/28] Add adaptive per-query fusion gated by vector-arm confidence The retrieval benchmark showed no single lexical_fusion_weight wins every corpus: SciFact wants a strong BM25 arm, NFCorpus a weak one, FiQA none. Rather than fix one weight, adaptive_fusion (off by default) scales the BM25 arm per query by how peaked the vector ranking is -- a top hit standing clear of the field downweights lexical toward zero, a flat ranking keeps it. adaptive_fusion_ margin sets the vector-similarity margin at which the lexical arm is fully silenced. lexical_fusion_weight becomes the ceiling the rule scales down from, so the default path is unchanged. (cherry picked from commit a0c2a96cb3b3880e62237ce44eae7ff76d6ed325) --- docs/usage.md | 2 ++ src/lilbee/app/settings_map.py | 12 ++++++++ src/lilbee/core/config/model.py | 10 +++++++ src/lilbee/data/store/core.py | 13 ++++++-- src/lilbee/data/store/fusion.py | 35 ++++++++++++++++++++++ tests/test_store.py | 11 +++++++ tests/test_store_fusion.py | 53 ++++++++++++++++++++++++++++++++- 7 files changed, 132 insertions(+), 4 deletions(-) diff --git a/docs/usage.md b/docs/usage.md index 7d2917ca1..e65a3515c 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -751,6 +751,8 @@ something feels off. | `LILBEE_TITLE_SEARCH` | `false` | Match queries against document titles as a third hybrid-search arm, so a query naming a document by title surfaces its chunks | | `LILBEE_TITLE_SEARCH_WEIGHT` | `0.5` | Title arm weight in rank fusion (1.0 = equal voice with the vector and text arms) | | `LILBEE_LEXICAL_FUSION_WEIGHT` | `1.0` | BM25 arm weight in rank fusion (1.0 = equal to the vector arm; lower it when a strong embedder should dominate a corpus where the lexical arm adds noise) | +| `LILBEE_ADAPTIVE_FUSION` | `false` | Scale the BM25 arm's weight per query by how confident the vector arm is (a peaked dense ranking downweights lexical, a flat one keeps it) instead of using a fixed `LILBEE_LEXICAL_FUSION_WEIGHT`. That weight becomes the ceiling the adaptive rule scales down from | +| `LILBEE_ADAPTIVE_FUSION_MARGIN` | `0.3` | Vector-similarity margin (top hit minus the field) at which adaptive fusion fully silences the BM25 arm; smaller values downweight lexical more aggressively | | `LILBEE_ADAPTIVE_THRESHOLD` | `false` | Widen the distance threshold step by step when too few results pass, on the vector-only fallback path (no effect on the default hybrid search) | | `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 | diff --git a/src/lilbee/app/settings_map.py b/src/lilbee/app/settings_map.py index 95fa5d58f..d3c4eb941 100644 --- a/src/lilbee/app/settings_map.py +++ b/src/lilbee/app/settings_map.py @@ -835,6 +835,18 @@ def get_default(key: str) -> object: group=SettingGroup.RETRIEVAL, help_text="BM25 arm weight in fusion (1.0 = equal to vector; lower to favor dense)", ), + "adaptive_fusion": SettingDef( + bool, + nullable=False, + group=SettingGroup.RETRIEVAL, + help_text="Scale the BM25 weight per query by vector-arm confidence, not a fixed value", + ), + "adaptive_fusion_margin": SettingDef( + float, + nullable=False, + group=SettingGroup.RETRIEVAL, + help_text="Vector-similarity margin at which adaptive fusion fully silences the BM25 arm", + ), "history_rewrite": SettingDef( bool, nullable=False, diff --git a/src/lilbee/core/config/model.py b/src/lilbee/core/config/model.py index 0ad011a9f..26e4e21f6 100644 --- a/src/lilbee/core/config/model.py +++ b/src/lilbee/core/config/model.py @@ -210,6 +210,16 @@ class Config(BaseSettings): # the retrieval benchmark, not guessed here. lexical_fusion_weight: float = ConfigField(default=1.0, ge=0.0, le=1.0, writable=True) + # Adaptive fusion: instead of a fixed lexical_fusion_weight, scale the BM25 + # arm per query by how confident the vector arm is (a peaked dense ranking + # downweights lexical, a flat one keeps it). The benchmark found no single + # fixed weight wins every corpus, so this gates the weight per query. Off by + # default; when on, lexical_fusion_weight is the ceiling the adaptive rule + # scales down from. adaptive_fusion_margin is the vector-similarity margin at + # which the lexical arm is fully silenced (smaller = more aggressive). + adaptive_fusion: bool = ConfigField(default=False, writable=True) + adaptive_fusion_margin: float = ConfigField(default=0.3, ge=0.0, le=2.0, writable=True) + # Chunk count at/above which sync builds an approximate (ANN) vector index # so search stays fast at millions of vectors. Below this, search uses exact # flat scan (faster and exact for small vaults). 0 disables the ANN index. diff --git a/src/lilbee/data/store/core.py b/src/lilbee/data/store/core.py index fa09b03f9..427fb1d53 100644 --- a/src/lilbee/data/store/core.py +++ b/src/lilbee/data/store/core.py @@ -28,7 +28,7 @@ from lilbee.core.security import validate_path_within from lilbee.runtime.lock import LOCK_TIMEOUT, write_lock -from .fusion import fuse_arms, normalized_bm25, vector_similarity +from .fusion import adaptive_lexical_weight, fuse_arms, normalized_bm25, vector_similarity from .lance_helpers import ( _chunk_type_predicate, _has_fts_index, @@ -712,11 +712,18 @@ def _hybrid_search( title_rows: list[SearchChunk] = [] if self._config.title_search: title_rows = self._title_arm(table, query_text, top_k, chunk_type) + vector_rows = self._vector_arm(table, query_vector, top_k, chunk_type) + lexical_weight = self._config.lexical_fusion_weight + if self._config.adaptive_fusion: + # Gate the lexical arm per query by how peaked the vector ranking is. + lexical_weight = adaptive_lexical_weight( + vector_rows, lexical_weight, self._config.adaptive_fusion_margin + ) fused = fuse_arms( - self._vector_arm(table, query_vector, top_k, chunk_type), + vector_rows, self._fts_arm(table, query_text, top_k, chunk_type), title_rows, - lexical_weight=self._config.lexical_fusion_weight, + lexical_weight=lexical_weight, title_weight=self._config.title_search_weight, ) fused = _drop_unsupported_far_rows(fused, max_distance) diff --git a/src/lilbee/data/store/fusion.py b/src/lilbee/data/store/fusion.py index 8804a4a3f..727b6819a 100644 --- a/src/lilbee/data/store/fusion.py +++ b/src/lilbee/data/store/fusion.py @@ -26,17 +26,52 @@ from __future__ import annotations +from statistics import fmean + from .types import SearchChunk # Standard RRF smoothing constant (Cormack, Clarke & Buettcher 2009). _RRF_K = 60 +# Adaptive fusion needs a top hit plus at least one field row to measure a margin. +_MIN_ROWS_FOR_MARGIN = 2 + def vector_similarity(distance: float) -> float: """Cosine distance to canonical [0, 1] similarity (distance spans [0, 2]).""" return max(0.0, min(1.0, 1.0 - distance)) +def adaptive_lexical_weight( + vector_rows: list[SearchChunk], base_weight: float, margin_scale: float +) -> float: + """Shrink the lexical arm's weight toward zero when the vector arm is + confident about this query. + + A *peaked* vector ranking -- a top hit standing well clear of the field -- + means the dense embedder already located the answer and the lexical arm + mostly adds term-match noise (the FiQA case in the retrieval benchmark). A + *flat* ranking means dense is unsure and BM25's exact-term matching is worth + trusting (NFCorpus, SciFact). The confidence signal is the margin between + the top similarity and the mean of the rest, divided by *margin_scale*: at or + above that margin the lexical arm is fully silenced, at zero margin it keeps + its full *base_weight*, and it scales linearly between. Returns *base_weight* + unchanged when there is nothing to measure (fewer than two scored rows) or + the feature is disabled (*margin_scale* <= 0). + """ + if margin_scale <= 0: + return base_weight + sims = sorted( + (vector_similarity(r.distance) for r in vector_rows if r.distance is not None), + reverse=True, + ) + if len(sims) < _MIN_ROWS_FOR_MARGIN: + return base_weight + margin = max(0.0, sims[0] - fmean(sims[1:])) + confidence = min(1.0, margin / margin_scale) + return base_weight * (1.0 - confidence) + + def normalized_bm25(scores: list[float]) -> list[float]: """Scale raw BM25 scores against the list maximum, into (0, 1]. diff --git a/tests/test_store.py b/tests/test_store.py index d7b877ec9..b520adcec 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -644,6 +644,17 @@ def test_hybrid_search_with_fts_index(self, store, test_config): scores = [r.score for r in results] assert all(0.0 <= s <= 1.0 for s in scores) + def test_adaptive_fusion_gates_the_hybrid_search(self, store, test_config): + """With adaptive_fusion on, hybrid search derives the lexical weight per + query from the vector arm instead of the fixed config value.""" + test_config.adaptive_fusion = True + store.add_chunks(_make_records()) + store.ensure_fts_index() + query_vec = [0.5] * test_config.embedding_dim + results = store.search(query_vec, top_k=3, query_text="chunk number") + assert len(results) > 0 + assert all(r.score is not None and 0.0 <= r.score <= 1.0 for r in results) + def test_fallback_to_vector_when_no_query_text(self, store, test_config): records = _make_records() store.add_chunks(records) diff --git a/tests/test_store_fusion.py b/tests/test_store_fusion.py index ac3c352dc..15762bf5b 100644 --- a/tests/test_store_fusion.py +++ b/tests/test_store_fusion.py @@ -3,7 +3,12 @@ import pytest from lilbee.data.store import SearchChunk -from lilbee.data.store.fusion import fuse_arms, normalized_bm25, vector_similarity +from lilbee.data.store.fusion import ( + adaptive_lexical_weight, + fuse_arms, + normalized_bm25, + vector_similarity, +) def _chunk(source, idx, *, distance=None, bm25=None, dim=4): @@ -68,6 +73,52 @@ def test_lexical_only_row_survives_with_score(self): assert lexical_row.score == pytest.approx(0.5) +class TestAdaptiveLexicalWeight: + """Per-query lexical weight gated by the vector arm's confidence.""" + + def test_peaked_dense_silences_lexical(self): + # top similarity 1.0 (distance 0), field ~0.3: a wide margin => weight ~0. + rows = [_chunk("a.md", 0, distance=0.0)] + [ + _chunk("b.md", i, distance=0.7) for i in range(1, 5) + ] + assert adaptive_lexical_weight(rows, 1.0, 0.3) == pytest.approx(0.0) + + def test_flat_dense_keeps_full_weight(self): + # every row equally similar: zero margin => the arm keeps base_weight. + rows = [_chunk("a.md", i, distance=0.5) for i in range(5)] + assert adaptive_lexical_weight(rows, 1.0, 0.3) == pytest.approx(1.0) + + def test_scales_linearly_between(self): + # top sim 0.6, field 0.3, margin 0.3, scale 0.6 => confidence 0.5 => half. + rows = [_chunk("a.md", 0, distance=0.4)] + [ + _chunk("b.md", i, distance=0.7) for i in range(1, 4) + ] + assert adaptive_lexical_weight(rows, 1.0, 0.6) == pytest.approx(0.5) + + def test_respects_base_weight(self): + rows = [_chunk("a.md", i, distance=0.5) for i in range(3)] + assert adaptive_lexical_weight(rows, 0.5, 0.3) == pytest.approx(0.5) + + def test_too_few_rows_returns_base(self): + assert adaptive_lexical_weight([_chunk("a.md", 0, distance=0.1)], 1.0, 0.3) == 1.0 + assert adaptive_lexical_weight([], 1.0, 0.3) == 1.0 + + def test_non_positive_margin_scale_disables(self): + rows = [_chunk("a.md", 0, distance=0.0), _chunk("b.md", 1, distance=0.9)] + assert adaptive_lexical_weight(rows, 1.0, 0.0) == 1.0 + + def test_ignores_rows_without_distance(self): + # lexical-only rows carry no distance; they must not enter the signal. + rows = [ + _chunk("a.md", 0, distance=0.0), + _chunk("b.md", 1, distance=0.6), + _chunk("c.md", 2, bm25=9.0), + ] + both = adaptive_lexical_weight(rows, 1.0, 0.4) + two = adaptive_lexical_weight(rows[:2], 1.0, 0.4) + assert both == pytest.approx(two) + + class TestLexicalFusionWeight: """The BM25 arm's fusion weight can be lowered so a strong dense arm dominates.""" From c58b1bd57b74c74c0ad15cdf89ecf944ae1d3596 Mon Sep 17 00:00:00 2001 From: Tobias Perelstein <5562156+tobocop2@users.noreply.github.com> Date: Sat, 18 Jul 2026 00:45:42 -0400 Subject: [PATCH 08/28] Turn adaptive fusion on by default The retrieval benchmark validated adaptive per-query fusion against the fixed lexical weight on BEIR SciFact, NFCorpus, and FiQA: at margin 0.15 it beats the fixed w=1.0 default on all three, keeps the significant NFCorpus and SciFact wins, and cuts the FiQA regression from -0.056 to -0.018 (the biggest gain on the one corpus a fixed lexical arm hurt most). Default adaptive_fusion=true, adaptive_fusion_margin=0.15; set adaptive_fusion=false to pin the fixed weight. --- docs/usage.md | 4 ++-- src/lilbee/core/config/model.py | 17 ++++++++++------- tests/test_store.py | 16 +++++++++++++--- 3 files changed, 25 insertions(+), 12 deletions(-) diff --git a/docs/usage.md b/docs/usage.md index e65a3515c..8e82aec69 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -751,8 +751,8 @@ something feels off. | `LILBEE_TITLE_SEARCH` | `false` | Match queries against document titles as a third hybrid-search arm, so a query naming a document by title surfaces its chunks | | `LILBEE_TITLE_SEARCH_WEIGHT` | `0.5` | Title arm weight in rank fusion (1.0 = equal voice with the vector and text arms) | | `LILBEE_LEXICAL_FUSION_WEIGHT` | `1.0` | BM25 arm weight in rank fusion (1.0 = equal to the vector arm; lower it when a strong embedder should dominate a corpus where the lexical arm adds noise) | -| `LILBEE_ADAPTIVE_FUSION` | `false` | Scale the BM25 arm's weight per query by how confident the vector arm is (a peaked dense ranking downweights lexical, a flat one keeps it) instead of using a fixed `LILBEE_LEXICAL_FUSION_WEIGHT`. That weight becomes the ceiling the adaptive rule scales down from | -| `LILBEE_ADAPTIVE_FUSION_MARGIN` | `0.3` | Vector-similarity margin (top hit minus the field) at which adaptive fusion fully silences the BM25 arm; smaller values downweight lexical more aggressively | +| `LILBEE_ADAPTIVE_FUSION` | `true` | Scale the BM25 arm's weight per query by how confident the vector arm is (a peaked dense ranking downweights lexical, a flat one keeps it) instead of using a fixed `LILBEE_LEXICAL_FUSION_WEIGHT`. That weight becomes the ceiling the adaptive rule scales down from. Set to `false` to pin the fixed weight | +| `LILBEE_ADAPTIVE_FUSION_MARGIN` | `0.15` | Vector-similarity margin (top hit minus the field) at which adaptive fusion fully silences the BM25 arm; smaller values downweight lexical more aggressively | | `LILBEE_ADAPTIVE_THRESHOLD` | `false` | Widen the distance threshold step by step when too few results pass, on the vector-only fallback path (no effect on the default hybrid search) | | `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 | diff --git a/src/lilbee/core/config/model.py b/src/lilbee/core/config/model.py index 26e4e21f6..5f8884cc9 100644 --- a/src/lilbee/core/config/model.py +++ b/src/lilbee/core/config/model.py @@ -212,13 +212,16 @@ class Config(BaseSettings): # Adaptive fusion: instead of a fixed lexical_fusion_weight, scale the BM25 # arm per query by how confident the vector arm is (a peaked dense ranking - # downweights lexical, a flat one keeps it). The benchmark found no single - # fixed weight wins every corpus, so this gates the weight per query. Off by - # default; when on, lexical_fusion_weight is the ceiling the adaptive rule - # scales down from. adaptive_fusion_margin is the vector-similarity margin at - # which the lexical arm is fully silenced (smaller = more aggressive). - adaptive_fusion: bool = ConfigField(default=False, writable=True) - adaptive_fusion_margin: float = ConfigField(default=0.3, ge=0.0, le=2.0, writable=True) + # downweights lexical, a flat one keeps it). On by default: the retrieval + # benchmark found no single fixed weight wins every corpus, and adaptive at + # margin 0.15 beat the fixed-weight default on all three BEIR sets tested + # (biggest gain on the corpus a fixed lexical arm hurt most). lexical_fusion_ + # weight is the ceiling the adaptive rule scales down from; adaptive_fusion_ + # margin is the vector-similarity margin at which the lexical arm is fully + # silenced (smaller = more aggressive). Set adaptive_fusion=false to pin the + # fixed weight instead. + adaptive_fusion: bool = ConfigField(default=True, writable=True) + adaptive_fusion_margin: float = ConfigField(default=0.15, ge=0.0, le=2.0, writable=True) # Chunk count at/above which sync builds an approximate (ANN) vector index # so search stays fast at millions of vectors. Below this, search uses exact diff --git a/tests/test_store.py b/tests/test_store.py index b520adcec..3dfcb6eed 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -645,9 +645,19 @@ def test_hybrid_search_with_fts_index(self, store, test_config): assert all(0.0 <= s <= 1.0 for s in scores) def test_adaptive_fusion_gates_the_hybrid_search(self, store, test_config): - """With adaptive_fusion on, hybrid search derives the lexical weight per - query from the vector arm instead of the fixed config value.""" - test_config.adaptive_fusion = True + """With adaptive_fusion on (the default), hybrid search derives the + lexical weight per query from the vector arm instead of the fixed value.""" + assert test_config.adaptive_fusion is True # shipped default + store.add_chunks(_make_records()) + store.ensure_fts_index() + query_vec = [0.5] * test_config.embedding_dim + results = store.search(query_vec, top_k=3, query_text="chunk number") + assert len(results) > 0 + assert all(r.score is not None and 0.0 <= r.score <= 1.0 for r in results) + + def test_fixed_fusion_when_adaptive_disabled(self, store, test_config): + """Opting out of adaptive fusion pins the fixed lexical_fusion_weight.""" + test_config.adaptive_fusion = False store.add_chunks(_make_records()) store.ensure_fts_index() query_vec = [0.5] * test_config.embedding_dim From 495c4daa1199158b3dd586c30af84e2158b05715 Mon Sep 17 00:00:00 2001 From: Tobias Perelstein <5562156+tobocop2@users.noreply.github.com> Date: Sat, 18 Jul 2026 01:50:09 -0400 Subject: [PATCH 09/28] Filter tables-of-contents and cover pages out of search results The UFO A/B (bb-5bjp) traced a real context-precision drop to document-structure chunks entering retrieval: the title arm matches cover/title pages (they carry the doc title) and table extraction indexes tables-of-contents, and neither ever answers a question. A conservative detector drops unambiguous TOCs (dot-leaders to page numbers) and cover pages (classification banner + short + shouting-case) from the search results, on the top_k*2 candidate buffer so real passages remain. On by default; filter_structural_chunks=false keeps them. Faithfulness was unaffected in the A/B (the relevant passages were already retrieved), so this recovers precision without touching answer quality. (cherry picked from commit 8ecbbf88e7c7490979e84e32b79c2b16c02ec5c7) --- docs/usage.md | 1 + src/lilbee/app/settings_map.py | 6 ++ src/lilbee/core/config/model.py | 6 ++ src/lilbee/retrieval/query/searcher.py | 7 +++ src/lilbee/retrieval/query/structural.py | 64 +++++++++++++++++++++ tests/test_query.py | 25 ++++++++ tests/test_structural.py | 72 ++++++++++++++++++++++++ 7 files changed, 181 insertions(+) create mode 100644 src/lilbee/retrieval/query/structural.py create mode 100644 tests/test_structural.py diff --git a/docs/usage.md b/docs/usage.md index 2c28f5f5e..71f5bdc99 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -813,6 +813,7 @@ something feels off. | `LILBEE_LEXICAL_FUSION_WEIGHT` | `1.0` | BM25 arm weight in rank fusion (1.0 = equal to the vector arm; lower it when a strong embedder should dominate a corpus where the lexical arm adds noise) | | `LILBEE_ADAPTIVE_FUSION` | `true` | Scale the BM25 arm's weight per query by how confident the vector arm is (a peaked dense ranking downweights lexical, a flat one keeps it) instead of using a fixed `LILBEE_LEXICAL_FUSION_WEIGHT`. That weight becomes the ceiling the adaptive rule scales down from. Set to `false` to pin the fixed weight | | `LILBEE_ADAPTIVE_FUSION_MARGIN` | `0.15` | Vector-similarity margin (top hit minus the field) at which adaptive fusion fully silences the BM25 arm; smaller values downweight lexical more aggressively | +| `LILBEE_FILTER_STRUCTURAL_CHUNKS` | `true` | Drop tables-of-contents and cover/title pages from search results. The title and table arms surface these document-structure chunks, which never answer a question and only dilute context precision. Set to `false` to keep them | | `LILBEE_ADAPTIVE_THRESHOLD` | `false` | Widen the distance threshold step by step when too few results pass, on the vector-only fallback path (no effect on the default hybrid search) | | `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 | diff --git a/src/lilbee/app/settings_map.py b/src/lilbee/app/settings_map.py index 3e0f64ed3..47b4486b4 100644 --- a/src/lilbee/app/settings_map.py +++ b/src/lilbee/app/settings_map.py @@ -874,6 +874,12 @@ def get_default(key: str) -> object: group=SettingGroup.RETRIEVAL, help_text="Vector-similarity margin at which adaptive fusion fully silences the BM25 arm", ), + "filter_structural_chunks": SettingDef( + bool, + nullable=False, + group=SettingGroup.RETRIEVAL, + help_text="Drop tables-of-contents and cover pages from search results", + ), "history_rewrite": SettingDef( bool, nullable=False, diff --git a/src/lilbee/core/config/model.py b/src/lilbee/core/config/model.py index e28ab7122..b0f999d83 100644 --- a/src/lilbee/core/config/model.py +++ b/src/lilbee/core/config/model.py @@ -223,6 +223,12 @@ class Config(BaseSettings): adaptive_fusion: bool = ConfigField(default=True, writable=True) adaptive_fusion_margin: float = ConfigField(default=0.15, ge=0.0, le=2.0, writable=True) + # Drop tables-of-contents and cover/title pages from search results. The + # title and table arms surface these document-structure chunks, which never + # answer a question and only dilute context precision. On by default; the + # detector is conservative (unambiguous TOCs and cover pages only). + filter_structural_chunks: bool = ConfigField(default=True, writable=True) + # Chunk count at/above which sync builds an approximate (ANN) vector index # so search stays fast at millions of vectors. Below this, search uses exact # flat scan (faster and exact for small vaults). 0 disables the ANN index. diff --git a/src/lilbee/retrieval/query/searcher.py b/src/lilbee/retrieval/query/searcher.py index 81f1916a8..2ee1f00a4 100644 --- a/src/lilbee/retrieval/query/searcher.py +++ b/src/lilbee/retrieval/query/searcher.py @@ -71,6 +71,7 @@ ) from lilbee.retrieval.query.memory import format_memory_block from lilbee.retrieval.query.neighbors import expand_neighbors +from lilbee.retrieval.query.structural import is_structural_chunk from lilbee.retrieval.query.tokenize import _idf_weights, _tokenize from lilbee.retrieval.reasoning import ( StreamToken, @@ -623,6 +624,12 @@ def search( # HTTP, and MCP copies of a bare distance cutoff dropped both-arm # rows the fusion layer deliberately keeps past max_distance. results = filter_results(results, self._config.max_distance) + # Drop tables-of-contents and cover pages: the title and table arms + # surface them, but they never answer a question and only dilute + # context precision (bb-pkn6). Filtered from the top_k*2 candidate + # buffer, so enough real passages remain for the downstream trim. + if self._config.filter_structural_chunks: + results = [r for r in results if not is_structural_chunk(r.chunk)] return results[: top_k * 2] def _condense_question(self, question: str, history: list[ChatMessage]) -> str: diff --git a/src/lilbee/retrieval/query/structural.py b/src/lilbee/retrieval/query/structural.py new file mode 100644 index 000000000..58bd24357 --- /dev/null +++ b/src/lilbee/retrieval/query/structural.py @@ -0,0 +1,64 @@ +"""Detect document-structure chunks that dilute retrieval precision. + +Tables of contents, cover/title pages, and bare reference lists carry a +document's title and section words, so the title arm and table extraction +surface them, but they never *answer* a substantive question. The UFO A/B +(bb-5bjp / bb-pkn6) traced a real context-precision drop to exactly these +chunks entering the retrieved set. The detector is deliberately conservative: +it protects recall/faithfulness by only flagging chunks that are unambiguously +structural, so a false negative (some noise slips through) is preferred to a +false positive (dropping real content). +""" + +from __future__ import annotations + +import re + +# A TOC line ends in dot leaders followed by a page number: "Geographic Trends ....... 9". +_TOC_LINE = re.compile(r"\.{3,}\s*\d{1,4}\s*$") + +# Classification banners head cover/title pages and running headers alike, so the +# banner alone is not enough -- it is combined with low prose density below. +_CLASSIFICATION = re.compile(r"\b(UNCLASSIFIED|CONFIDENTIAL|SECRET|FOR OFFICIAL USE ONLY|FOUO)\b") + +# A chunk needs at least this many non-empty lines before the TOC ratio means anything. +_MIN_TOC_LINES = 3 +_TOC_RATIO = 0.30 + +# Cover/title-page gates: short, few sentences, shouting-case dominated. +_COVER_MAX_WORDS = 120 +_COVER_MAX_SENTENCES = 3 +_COVER_CAPS_RATIO = 0.30 + + +def _is_toc(nonempty: list[str]) -> bool: + """A table of contents: several dot-leader-to-page-number lines.""" + if len(nonempty) < _MIN_TOC_LINES: + return False + hits = sum(1 for line in nonempty if _TOC_LINE.search(line)) + return hits >= _MIN_TOC_LINES and hits / len(nonempty) >= _TOC_RATIO + + +def _is_cover_page(text: str) -> bool: + """A cover/title page: short, almost no sentences, shouting-case dominated, + and carrying a classification banner. Real prose has many sentence stops and + fails the sentence gate, so it is never flagged.""" + if not _CLASSIFICATION.search(text): + return False + words = text.split() + if not words or len(words) > _COVER_MAX_WORDS: + return False + sentences = text.count(".") + text.count("!") + text.count("?") + if sentences > _COVER_MAX_SENTENCES: + return False + caps = sum(1 for w in words if len(w) > 1 and w.isupper()) + return caps / len(words) >= _COVER_CAPS_RATIO + + +def is_structural_chunk(text: str) -> bool: + """True when *text* is a table of contents or a cover/title page -- a + document-structure chunk that should not compete as an answer passage.""" + if not text or not text.strip(): + return False + nonempty = [line for line in text.splitlines() if line.strip()] + return _is_toc(nonempty) or _is_cover_page(text) diff --git a/tests/test_query.py b/tests/test_query.py index 14fb47d78..bd680618e 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -501,6 +501,31 @@ def test_far_vector_only_rows_filtered_at_search(self, mock_svc): results = get_services().searcher.search("q") assert [r.source for r in results] == ["a.md"] + def test_structural_chunks_filtered_from_results(self, mock_svc): + """A TOC the title/table arms surfaced is dropped so it does not dilute + context precision, while the real answer passage survives (bb-pkn6).""" + toc = _make_result( + source="toc.pdf", + distance=0.2, + chunk="A. Summary ......... 1\nB. Intro ......... 3\nC. Trends ......... 9\n", + ) + real = _make_result(source="body.pdf", distance=0.2, chunk="Real answer prose.") + mock_svc.store.search.return_value = [toc, real] + results = get_services().searcher.search("q") + assert [r.source for r in results] == ["body.pdf"] + + def test_structural_filter_off_keeps_toc(self, mock_svc): + """Opting out (filter_structural_chunks=false) keeps every retrieved row.""" + cfg.filter_structural_chunks = False + toc = _make_result( + source="toc.pdf", + distance=0.2, + chunk="A. Summary ......... 1\nB. Intro ......... 3\nC. Trends ......... 9\n", + ) + mock_svc.store.search.return_value = [toc] + results = get_services().searcher.search("q") + assert [r.source for r in results] == ["toc.pdf"] + def test_far_row_with_lexical_support_survives_search(self, mock_svc): """A both-arm row keeps its standing past max_distance: dropping it on vector distance alone would re-bury exactly the identifier hits diff --git a/tests/test_structural.py b/tests/test_structural.py new file mode 100644 index 000000000..bb58f7ed6 --- /dev/null +++ b/tests/test_structural.py @@ -0,0 +1,72 @@ +"""Tests for the document-structure (TOC / cover-page) chunk detector.""" + +from lilbee.retrieval.query.structural import is_structural_chunk + +# A real table of contents, as seen diluting precision in the UFO A/B (bb-pkn6). +TOC = """Contents +A. Executive Summary ................................................. 1 +B. Introduction ..................................................... 3 +C. Geographic Trends ................................................ 9 +D. Notable Trends Regarding Propulsion .............................. 10 +E. Flight Safety Issues ............................................. 11 +F. UAS Observations Reported ........................................ 13 +""" + +# A real DoD cover/title page. +COVER = """UNCLASSIFIED +THE DEPARTMENT OF DEFENSE +ALL-DOMAIN ANOMALY RESOLUTION OFFICE +Fiscal Year 2024 Consolidated Annual Report on +Unidentified Anomalous Phenomena +Information Cut Off Date: 1 JUNE 2024 +UNCLASSIFIED +""" + +# Real prose that must NOT be flagged. +PROSE = ( + "During the reporting period, AARO received no reports indicating UAP sightings " + "have been associated with any adverse health effects. However, many reports from " + "military witnesses described transient effects. The office continues to evaluate " + "each case against a standardized methodology, and it has resolved the majority of " + "reported incidents as ordinary objects or sensor artifacts." +) + +# Prose that happens to reference a page and carry a classification header. +PROSE_WITH_HEADER = ( + "UNCLASSIFIED. As detailed on page 12, the assessment concluded that the object was " + "a commercial aircraft. The radar track and the electro-optical imagery were " + "consistent, and the case was closed. No anomalous performance was observed at any " + "point during the encounter, which lasted several minutes." +) + + +class TestIsStructuralChunk: + def test_table_of_contents_is_structural(self): + assert is_structural_chunk(TOC) is True + + def test_cover_page_is_structural(self): + assert is_structural_chunk(COVER) is True + + def test_real_prose_is_not_structural(self): + assert is_structural_chunk(PROSE) is False + + def test_prose_with_classification_header_and_page_ref_is_not_structural(self): + # Many sentences => fails the cover-page gate; one "page 12" => not a TOC. + assert is_structural_chunk(PROSE_WITH_HEADER) is False + + def test_empty_is_not_structural(self): + assert is_structural_chunk("") is False + assert is_structural_chunk(" \n ") is False + + def test_single_dot_leader_line_is_not_a_toc(self): + assert is_structural_chunk("See the appendix ......... 42") is False + + def test_short_all_caps_without_classification_is_not_a_cover(self): + # A shouting heading with no classification banner is left alone. + assert is_structural_chunk("NOTABLE TRENDS REGARDING PROPULSION AND FLIGHT") is False + + def test_long_classified_body_is_not_a_cover(self): + # A real document body that opens with a classification banner but runs + # long is content, not a cover page: the word-count gate protects it. + body = "UNCLASSIFIED " + " ".join(f"finding{i} detail" for i in range(80)) + assert is_structural_chunk(body) is False From 0c598d6405846b110cdd1c4f66680d5435b4fd6b Mon Sep 17 00:00:00 2001 From: Tobias Perelstein <5562156+tobocop2@users.noreply.github.com> Date: Sat, 18 Jul 2026 19:00:22 -0400 Subject: [PATCH 10/28] Default the structural-chunk filter off and stop it dropping needed pages A UFO-corpus A/B (bb-lenb) found filter_structural_chunks net-negative: its cover-page heuristic also fired on short, classification-banner government body pages, dropping content the answer needed and hurting faithfulness (a timeline question went from 1.00 to 0.10). Default it off pending per-corpus evidence, never drop a query-matched (BM25) or top-ranked page whatever its shape, and tighten the cover-page gates (<=60 words, <=1 sentence) so ordinary short classified body pages no longer trip it. Regression tests cover all three. --- src/lilbee/core/config/model.py | 13 ++++-- src/lilbee/retrieval/query/searcher.py | 10 ++++- src/lilbee/retrieval/query/structural.py | 10 +++-- tests/test_query.py | 51 ++++++++++++++---------- tests/test_structural.py | 11 +++++ 5 files changed, 67 insertions(+), 28 deletions(-) diff --git a/src/lilbee/core/config/model.py b/src/lilbee/core/config/model.py index b0f999d83..b3068cd88 100644 --- a/src/lilbee/core/config/model.py +++ b/src/lilbee/core/config/model.py @@ -224,10 +224,15 @@ class Config(BaseSettings): adaptive_fusion_margin: float = ConfigField(default=0.15, ge=0.0, le=2.0, writable=True) # Drop tables-of-contents and cover/title pages from search results. The - # title and table arms surface these document-structure chunks, which never - # answer a question and only dilute context precision. On by default; the - # detector is conservative (unambiguous TOCs and cover pages only). - filter_structural_chunks: bool = ConfigField(default=True, writable=True) + # title and table arms surface these document-structure chunks, which usually + # do not answer a question. OFF by default: a UFO-corpus A/B (bb-lenb) found + # the filter net-negative -- its cover-page heuristic also fires on short, + # classification-banner government body pages, dropping content the answer + # needed and hurting faithfulness. When on, a query-matched or top-ranked + # page is never dropped (searcher.search), so the removal is limited to + # structural chunks the query did not actually hit. Re-validate per corpus + # before turning it on. + filter_structural_chunks: bool = ConfigField(default=False, writable=True) # Chunk count at/above which sync builds an approximate (ANN) vector index # so search stays fast at millions of vectors. Below this, search uses exact diff --git a/src/lilbee/retrieval/query/searcher.py b/src/lilbee/retrieval/query/searcher.py index 2ee1f00a4..f16aaa3ee 100644 --- a/src/lilbee/retrieval/query/searcher.py +++ b/src/lilbee/retrieval/query/searcher.py @@ -629,7 +629,15 @@ def search( # context precision (bb-pkn6). Filtered from the top_k*2 candidate # buffer, so enough real passages remain for the downstream trim. if self._config.filter_structural_chunks: - results = [r for r in results if not is_structural_chunk(r.chunk)] + # Never drop a page the query actually hit: a lexical (BM25) match or + # the top-ranked result is content the answer may need, whatever its + # shape, so only genuinely-unhit structural chunks are removed + # (bb-lenb: the filter was dropping needed body pages). + results = [ + r + for i, r in enumerate(results) + if r.bm25_score is not None or i == 0 or not is_structural_chunk(r.chunk) + ] return results[: top_k * 2] def _condense_question(self, question: str, history: list[ChatMessage]) -> str: diff --git a/src/lilbee/retrieval/query/structural.py b/src/lilbee/retrieval/query/structural.py index 58bd24357..e01e50b47 100644 --- a/src/lilbee/retrieval/query/structural.py +++ b/src/lilbee/retrieval/query/structural.py @@ -25,9 +25,13 @@ _MIN_TOC_LINES = 3 _TOC_RATIO = 0.30 -# Cover/title-page gates: short, few sentences, shouting-case dominated. -_COVER_MAX_WORDS = 120 -_COVER_MAX_SENTENCES = 3 +# Cover/title-page gates: a title page is very short with essentially no prose. +# These are deliberately tight -- a UFO-corpus A/B (bb-lenb) showed looser gates +# firing on short, classification-banner government BODY pages and dropping the +# content the answer needed, so a real body page's word count or its first full +# sentence must take it out of scope. +_COVER_MAX_WORDS = 60 +_COVER_MAX_SENTENCES = 1 _COVER_CAPS_RATIO = 0.30 diff --git a/tests/test_query.py b/tests/test_query.py index bd680618e..24bbb9a25 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -501,30 +501,41 @@ def test_far_vector_only_rows_filtered_at_search(self, mock_svc): results = get_services().searcher.search("q") assert [r.source for r in results] == ["a.md"] - def test_structural_chunks_filtered_from_results(self, mock_svc): - """A TOC the title/table arms surfaced is dropped so it does not dilute - context precision, while the real answer passage survives (bb-pkn6).""" - toc = _make_result( - source="toc.pdf", - distance=0.2, - chunk="A. Summary ......... 1\nB. Intro ......... 3\nC. Trends ......... 9\n", - ) - real = _make_result(source="body.pdf", distance=0.2, chunk="Real answer prose.") - mock_svc.store.search.return_value = [toc, real] + _TOC_CHUNK = "A. Summary ......... 1\nB. Intro ......... 3\nC. Trends ......... 9\n" + + def test_structural_chunk_the_query_missed_is_dropped(self, mock_svc): + """With the filter on, a lower-ranked TOC the query did not hit is + dropped, while the real answer passage survives (bb-pkn6).""" + cfg.filter_structural_chunks = True + real = _make_result(source="body.pdf", distance=0.1, chunk="Real answer prose.") + toc = _make_result(source="toc.pdf", distance=0.4, chunk=self._TOC_CHUNK) + mock_svc.store.search.return_value = [real, toc] results = get_services().searcher.search("q") assert [r.source for r in results] == ["body.pdf"] - def test_structural_filter_off_keeps_toc(self, mock_svc): - """Opting out (filter_structural_chunks=false) keeps every retrieved row.""" - cfg.filter_structural_chunks = False - toc = _make_result( - source="toc.pdf", - distance=0.2, - chunk="A. Summary ......... 1\nB. Intro ......... 3\nC. Trends ......... 9\n", - ) + def test_structural_filter_keeps_query_matched_page(self, mock_svc): + """A page the query lexically hit is never dropped, whatever its shape: + it is content the answer may need (bb-lenb).""" + cfg.filter_structural_chunks = True + real = _make_result(source="body.pdf", distance=0.1, chunk="Real answer prose.") + hit = _make_result(source="toc.pdf", distance=0.4, bm25_score=9.0, chunk=self._TOC_CHUNK) + mock_svc.store.search.return_value = [real, hit] + assert "toc.pdf" in [r.source for r in get_services().searcher.search("q")] + + def test_structural_filter_keeps_top_hit(self, mock_svc): + """The top-ranked result is never dropped, whatever its shape (bb-lenb).""" + cfg.filter_structural_chunks = True + toc = _make_result(source="toc.pdf", distance=0.05, chunk=self._TOC_CHUNK) mock_svc.store.search.return_value = [toc] - results = get_services().searcher.search("q") - assert [r.source for r in results] == ["toc.pdf"] + assert [r.source for r in get_services().searcher.search("q")] == ["toc.pdf"] + + def test_structural_filter_off_by_default_keeps_toc(self, mock_svc): + """Off by default (bb-lenb: net-negative on the UFO corpus).""" + assert cfg.filter_structural_chunks is False + toc = _make_result(source="toc.pdf", distance=0.4, chunk=self._TOC_CHUNK) + real = _make_result(source="body.pdf", distance=0.1, chunk="Real answer prose.") + mock_svc.store.search.return_value = [real, toc] + assert "toc.pdf" in [r.source for r in get_services().searcher.search("q")] def test_far_row_with_lexical_support_survives_search(self, mock_svc): """A both-arm row keeps its standing past max_distance: dropping it diff --git a/tests/test_structural.py b/tests/test_structural.py index bb58f7ed6..4d36d02b1 100644 --- a/tests/test_structural.py +++ b/tests/test_structural.py @@ -65,6 +65,17 @@ def test_short_all_caps_without_classification_is_not_a_cover(self): # A shouting heading with no classification banner is left alone. assert is_structural_chunk("NOTABLE TRENDS REGARDING PROPULSION AND FLIGHT") is False + def test_short_classified_body_page_is_not_a_cover(self): + # A short government body page carries a classification banner and caps + # but real content; its full sentences must keep it out of scope so the + # answer does not lose the page it needs (bb-lenb, the rag-1 failure). + body = ( + "UNCLASSIFIED. The AARO assessment concluded the object was a " + "commercial aircraft, and the case was resolved. RADAR and EO/IR " + "data were CONSISTENT across the entire track." + ) + assert is_structural_chunk(body) is False + def test_long_classified_body_is_not_a_cover(self): # A real document body that opens with a classification banner but runs # long is content, not a cover page: the word-count gate protects it. From 3eeb98d90cdc5ef182d2c84cb7499694976a8db3 Mon Sep 17 00:00:00 2001 From: Tobias Perelstein <5562156+tobocop2@users.noreply.github.com> Date: Sun, 19 Jul 2026 01:09:22 -0400 Subject: [PATCH 11/28] Build FTS indexes without token positions to stop the optimize() overflow with_position=True made the chunk index's trigram position lists overflow LanceDB's list encoding on a large corpus, so optimize() failed every serve start and BM25 was served un-optimized all session (bb-rqr8, confirmed on the failing index: rebuilding the chunk index positionless makes optimize() succeed). Positions only serve exact-phrase queries, which lilbee never issues, so build both FTS indexes positionless and sanitize the double quotes that would make LanceDB parse a query as a phrase (_sanitize_fts_query), applied at every FTS query site so a positionless index never errors on a quoted query. --- src/lilbee/data/store/core.py | 42 ++++++++++++++++++++++++++++------- tests/test_store.py | 23 ++++++++++++------- 2 files changed, 49 insertions(+), 16 deletions(-) diff --git a/src/lilbee/data/store/core.py b/src/lilbee/data/store/core.py index 9d5199663..850f8e2c0 100644 --- a/src/lilbee/data/store/core.py +++ b/src/lilbee/data/store/core.py @@ -106,6 +106,19 @@ def _drop_unsupported_far_rows( _MAX_THRESHOLD = 1.0 _MAX_FILTER_ITERATIONS = 20 # safety cap to prevent runaway loops + +def _sanitize_fts_query(text: str) -> str: + """Strip the double quotes that make LanceDB parse an FTS query as a phrase. + + lilbee's FTS indexes carry no token positions (bb-rqr8: positions overflow + LanceDB's list encoding on a large corpus), so a phrase query would error. + Replacing quotes with spaces turns a quoted span into plain terms -- which + is what every lilbee FTS call wants anyway, since none needs exact-phrase + matching. + """ + return text.replace('"', " ") + + # Vector ANN index. IVF_PQ compresses vectors so search scales to millions; # refine_factor re-ranks the PQ candidates against full vectors to recover recall. _VECTOR_METRIC = "cosine" @@ -444,10 +457,15 @@ def ensure_fts_index(self) -> None: exc_info=True, ) else: - # with_position lets the lexical arm serve phrase queries; - # raw user query text can reach LanceDB as a phrase, which - # errors on a positionless index. - table.create_fts_index("chunk", replace=False, with_position=True) + # No token positions: with_position=True makes the chunk + # index's trigram position lists overflow LanceDB's list + # encoding on optimize() for a large corpus, which then + # serves BM25 un-optimized all session (bb-rqr8). Positions + # only serve exact-phrase queries, which lilbee never issues + # -- query text is sanitized of the quotes that would make + # LanceDB parse it as a phrase (see _sanitize_fts_query), so + # a positionless index never errors on one. + table.create_fts_index("chunk", replace=False, with_position=False) self._fts_ready = True log.debug("FTS index created on '%s'", CHUNKS_TABLE) self._ensure_title_fts_unlocked(table) @@ -463,7 +481,9 @@ def _ensure_title_fts_unlocked(self, table: lancedb.table.Table) -> None: if _TITLE_COLUMN not in table.schema.names or _has_fts_index(table, _TITLE_COLUMN): return try: - table.create_fts_index(_TITLE_COLUMN, replace=False, with_position=True) + # Positionless for the same reason as the chunk index (bb-rqr8); + # queries are sanitized of phrase quotes so it never errors. + table.create_fts_index(_TITLE_COLUMN, replace=False, with_position=False) log.debug("Title FTS index created on '%s'", CHUNKS_TABLE) except Exception: log.debug("Title FTS index create failed", exc_info=True) @@ -548,7 +568,9 @@ def bm25_probe( # Pin the chunk column: once a title FTS index exists, an unpinned # FTS query searches every indexed column, which would silently # widen the probe's semantics. - query = table.search(query_text, query_type="fts", fts_columns="chunk") + query = table.search( + _sanitize_fts_query(query_text), query_type="fts", fts_columns="chunk" + ) if chunk_type: query = query.where(_chunk_type_predicate(chunk_type)) rows = query.limit(top_k).to_list() @@ -657,7 +679,9 @@ def _fts_arm( query searches every indexed column, which would double-count title tokens this arm's sibling already scores. """ - query = table.search(query_text, query_type="fts", fts_columns="chunk").limit(limit) + query = table.search( + _sanitize_fts_query(query_text), query_type="fts", fts_columns="chunk" + ).limit(limit) if chunk_type: query = query.where(_chunk_type_predicate(chunk_type)) return [SearchChunk(**r) for r in query.to_list()] @@ -676,7 +700,9 @@ def _title_arm( """ if not _has_fts_index(table, _TITLE_COLUMN): return [] - query = table.search(query_text, query_type="fts", fts_columns=_TITLE_COLUMN).limit(limit) + query = table.search( + _sanitize_fts_query(query_text), query_type="fts", fts_columns=_TITLE_COLUMN + ).limit(limit) if chunk_type: query = query.where(_chunk_type_predicate(chunk_type)) return [SearchChunk(**r) for r in query.to_list()] diff --git a/tests/test_store.py b/tests/test_store.py index dfccf9fd2..598340f41 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -192,16 +192,23 @@ def test_first_call_creates_without_replace(self, store): assert [c.args[0] for c in create_spy.call_args_list] == ["chunk", "title"] # Verify replace was NOT True (would defeat the purpose of incremental) assert all(c.kwargs.get("replace") is False for c in create_spy.call_args_list) - # Both indexes carry token positions so the lexical arm can serve - # phrase queries; without this LanceDB raises on any phrase. - assert all(c.kwargs.get("with_position") is True for c in create_spy.call_args_list) + # Both indexes are positionless: with_position=True overflows LanceDB's + # list encoding on a large corpus (bb-rqr8), and no lilbee query needs + # exact-phrase matching. + assert all(c.kwargs.get("with_position") is False for c in create_spy.call_args_list) - def test_fts_phrase_query_does_not_fail(self, store): - """A multi-word phrase query must return matches, not raise on a missing index. + def test_sanitize_fts_query_replaces_quotes_with_spaces(self): + from lilbee.data.store.core import _sanitize_fts_query - lilbee passes raw user query text to the lexical arm, so a quoted phrase - reaches LanceDB as a phrase query. Without positional indexing that raises - (and bm25_probe swallows it, returning nothing); with positions it matches. + assert _sanitize_fts_query('"exact phrase" and terms') == " exact phrase and terms" + assert _sanitize_fts_query("plain terms") == "plain terms" + + def test_fts_quoted_query_is_sanitized_not_a_phrase(self, store): + """A quoted query must return term matches, not raise on the positionless index. + + The chunk index carries no token positions (bb-rqr8), so a phrase query + would error. bm25_probe strips the quotes (_sanitize_fts_query), turning + the quoted span into plain terms that still match. """ store.add_chunks(_make_records()) store.ensure_fts_index() From 1ef3bdc7e1614deaa6884770cc988f3f0ef71074 Mon Sep 17 00:00:00 2001 From: Tobias Perelstein <5562156+tobocop2@users.noreply.github.com> Date: Sun, 19 Jul 2026 01:35:15 -0400 Subject: [PATCH 12/28] Index the columns lilbee filters by (source, chunk_type) source and chunk_type predicates already run as prefilters, but on unindexed columns each is a full-table scan. Build a BTree on the high-cardinality source (known-item lookup and per-source fetches) and a Bitmap on the low-cardinality chunk_type (raw/wiki/table scope) at ingest, next to the FTS and vector index builds, so those filters become indexed lookups. Purely a speedup: the query predicates are unchanged and stores without the indexes just run unindexed until re-ingested. --- src/lilbee/cli/commands/ingest_sync.py | 1 + src/lilbee/data/ingest/pipeline.py | 1 + src/lilbee/data/store/core.py | 26 +++++++++++++++ src/lilbee/data/store/lance_helpers.py | 13 ++++++++ tests/test_store.py | 44 ++++++++++++++++++++++++++ 5 files changed, 85 insertions(+) diff --git a/src/lilbee/cli/commands/ingest_sync.py b/src/lilbee/cli/commands/ingest_sync.py index e827106bf..b85564044 100644 --- a/src/lilbee/cli/commands/ingest_sync.py +++ b/src/lilbee/cli/commands/ingest_sync.py @@ -400,6 +400,7 @@ def index( apply_overrides(data_dir=data_dir, use_global=use_global) store = get_services().store store.ensure_fts_index() + store.ensure_scalar_indexes() built = store.ensure_vector_index(force=True) if cfg.json_mode: json_output({"command": "index", "vector_index": built}) diff --git a/src/lilbee/data/ingest/pipeline.py b/src/lilbee/data/ingest/pipeline.py index 0a1833d9a..df702b4f5 100644 --- a/src/lilbee/data/ingest/pipeline.py +++ b/src/lilbee/data/ingest/pipeline.py @@ -538,6 +538,7 @@ async def sync( if files_to_process or removed: _store.ensure_fts_index() + _store.ensure_scalar_indexes() _store.ensure_vector_index() _store.optimize_sources() await _rebuild_concept_clusters() diff --git a/src/lilbee/data/store/core.py b/src/lilbee/data/store/core.py index 850f8e2c0..018d04e0e 100644 --- a/src/lilbee/data/store/core.py +++ b/src/lilbee/data/store/core.py @@ -32,6 +32,7 @@ from .lance_helpers import ( _chunk_type_predicate, _has_fts_index, + _has_scalar_index, _has_vector_index, _safe_delete_unlocked, _sources_search_filter, @@ -488,6 +489,31 @@ def _ensure_title_fts_unlocked(self, table: lancedb.table.Table) -> None: except Exception: log.debug("Title FTS index create failed", exc_info=True) + def ensure_scalar_indexes(self) -> None: + """Build scalar indexes on the columns lilbee filters by. + + ``source`` and ``chunk_type`` predicates run as prefilters (LanceDB's + default), but without an index each is a full-table scan. A BTree on the + high-cardinality ``source`` (known-item lookup and per-source fetches) + and a Bitmap on the low-cardinality ``chunk_type`` (raw/wiki/table scope) + turn those into indexed lookups. Missing columns and empty tables are + skipped; ``optimize()`` folds later rows into every index, this one too. + """ + with self._write_lock(): + table = self.open_table(CHUNKS_TABLE) + if table is None: + return + names = table.schema.names + try: + if "source" in names and not _has_scalar_index(table, "source"): + table.create_scalar_index("source", index_type="BTREE", replace=False) + log.debug("Scalar (BTree) index created on 'source'") + if "chunk_type" in names and not _has_scalar_index(table, "chunk_type"): + table.create_scalar_index("chunk_type", index_type="BITMAP", replace=False) + log.debug("Scalar (Bitmap) index created on 'chunk_type'") + except Exception: + log.debug("Scalar index ensure failed (empty table?)", exc_info=True) + def ensure_vector_index(self, *, force: bool = False) -> bool: """Build or refresh the ANN vector index when the corpus is large enough. diff --git a/src/lilbee/data/store/lance_helpers.py b/src/lilbee/data/store/lance_helpers.py index a4b3a69a5..265e3ace3 100644 --- a/src/lilbee/data/store/lance_helpers.py +++ b/src/lilbee/data/store/lance_helpers.py @@ -139,6 +139,19 @@ def _has_fts_index(table: lancedb.table.Table, column: str = "chunk") -> bool: return False +def _has_scalar_index(table: lancedb.table.Table, column: str) -> bool: + """Return True when a scalar index on *column* already exists. + + lilbee only builds scalar indexes on ``source`` and ``chunk_type``, and + never an FTS or vector index on those columns, so any index touching the + column is the scalar one. + """ + try: + return any(column in idx.columns for idx in table.list_indices()) + except Exception: + return False + + def _has_vector_index(table: lancedb.table.Table) -> bool: """Return True when an ANN index on the vector column already exists. diff --git a/tests/test_store.py b/tests/test_store.py index 598340f41..ffad94830 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -305,6 +305,50 @@ def _make_indexable_records(n, dim): ] +class TestEnsureScalarIndexes: + """source and chunk_type get scalar indexes so their prefilters are lookups.""" + + def test_creates_btree_on_source_and_bitmap_on_chunk_type(self, store): + store.add_chunks(_make_records()) + table = store.open_table("chunks") + assert table is not None + with mock.patch.object(type(table), "create_scalar_index") as spy: + store.ensure_scalar_indexes() + assert [c.args[0] for c in spy.call_args_list] == ["source", "chunk_type"] + kinds = {c.args[0]: c.kwargs.get("index_type") for c in spy.call_args_list} + assert kinds == {"source": "BTREE", "chunk_type": "BITMAP"} + assert all(c.kwargs.get("replace") is False for c in spy.call_args_list) + + def test_idempotent_once_the_indexes_exist(self, store): + store.add_chunks(_make_records()) + store.ensure_scalar_indexes() # builds them for real + table = store.open_table("chunks") + with mock.patch.object(type(table), "create_scalar_index") as spy: + store.ensure_scalar_indexes() + spy.assert_not_called() + + def test_handles_exception_gracefully(self, store): + store.add_chunks(_make_records()) + table = store.open_table("chunks") + assert table is not None + with mock.patch.object( + type(table), "create_scalar_index", side_effect=RuntimeError("boom") + ): + store.ensure_scalar_indexes() # must not raise + + def test_noop_when_no_table(self, store): + store.ensure_scalar_indexes() # empty store, no chunks table yet + + def test_has_scalar_index_is_false_when_listing_raises(self, store): + from lilbee.data.store.lance_helpers import _has_scalar_index + + store.add_chunks(_make_records()) + table = store.open_table("chunks") + assert table is not None + with mock.patch.object(type(table), "list_indices", side_effect=RuntimeError("boom")): + assert _has_scalar_index(table, "source") is False + + class TestEnsureVectorIndex: """Small vaults stay on exact flat search; large ones get an ANN index.""" From bc65756e525838f8ceea1e9b4ca3a53e4402f07a Mon Sep 17 00:00:00 2001 From: Tobias Perelstein <5562156+tobocop2@users.noreply.github.com> Date: Sun, 19 Jul 2026 16:19:58 -0400 Subject: [PATCH 13/28] Address PR #557 audit findings: accuracy, conventions, test quality Remediates the bb-teg7 audit of the retrieval-consolidated branch. Code / behavior: - Consolidate the three duplicated FTS query sites into one _lexical_rows helper built on LanceDB's MatchQuery, deleting the hand-rolled quote sanitizer (verified identical rows and BM25 scores through both paths). - fuse_arms: skip a zero-weight arm instead of injecting zero-score rows that would still claim the distance and structural-filter exemptions. Conventions: - Remove corpus-identifying "UFO" references and bead IDs from src/ and tests/; rewrite the structural-detector fixtures corpus-neutrally. - Cut narrative and benchmark-storytelling comments to one-line invariants. Docs / docstrings: - _hybrid_search, fusion, and dedup docstrings, config comments, and docs/architecture.md + usage.md now describe the weighted N-arm adaptive fusion and the per-query score caveat, not the old fixed-0.5 two-arm story. - usage.md / settings_map: structural filter defaults off with honest classification-banner scope; adaptive_fusion_margin=0 documented as disabling adaptation. - searcher and structural comments now match the actual exemption behavior. Tests: - Move four general fuse_arms tests back into TestFuseArms (they had been reparented into TestLexicalFusionWeight). - Adaptive / fixed / title tests now spy on fuse_arms and adaptive_lexical_weight and assert the config weights and margin reach fusion, instead of only checking the score range. make check green: 9663 passed, 100% coverage. --- docs/architecture.md | 15 +-- docs/usage.md | 4 +- src/lilbee/app/settings_map.py | 2 +- src/lilbee/core/config/model.py | 43 ++++----- src/lilbee/data/ingest/types.py | 3 +- src/lilbee/data/store/core.py | 116 ++++++++++------------- src/lilbee/data/store/fusion.py | 29 ++++-- src/lilbee/retrieval/query/dedup.py | 11 ++- src/lilbee/retrieval/query/searcher.py | 15 ++- src/lilbee/retrieval/query/structural.py | 21 ++-- tests/test_query.py | 8 +- tests/test_store.py | 74 +++++++++++---- tests/test_store_fusion.py | 72 +++++++------- tests/test_structural.py | 50 +++++----- 14 files changed, 242 insertions(+), 221 deletions(-) diff --git a/docs/architecture.md b/docs/architecture.md index fce977469..4ddd3ac9e 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -641,7 +641,7 @@ flowchart TD SM -->|No prefix| PROBE[BM25 Confidence Probe] PROBE --> CONF{Confident AND separated?} - CONF -->|Yes| DUAL[Dual-Arm Retrieval] + CONF -->|Yes| DUAL[Hybrid Retrieval] CONF -->|No| EXPAND[LLM Query Expansion] EXPAND --> GEXP[+ Graph Expansion] @@ -677,19 +677,20 @@ flowchart TD #### History Condensation **On by default** (`LILBEE_HISTORY_REWRITE`). A follow-up question is condensed into a standalone retrieval query using the recent chat history (one small LLM call, skipped when there is no history). Retrieval sees only the query text: without this, "what about his brother?" is embedded and BM25-matched with its pronouns. The user's original wording still reaches the answering prompt. -#### Hybrid Search (Dual-Arm Reciprocal-Rank Fusion) -**Always on.** Retrieves the BM25 arm (LanceDB FTS) and the vector arm independently, each fetched exactly `top_k` deep, then fuses by reciprocal rank into one canonical [0, 1] score: +#### Hybrid Search (Weighted Reciprocal-Rank Fusion) +**Always on.** Retrieves a vector arm and a BM25 chunk arm (LanceDB FTS) independently, each fetched exactly `top_k` deep, then fuses by weighted reciprocal rank into one [0, 1] score. An optional third BM25 arm over document titles joins when `LILBEE_TITLE_SEARCH` is on: ``` -score = (rank_weight(vector_rank) + rank_weight(bm25_rank)) / 2, rank_weight(r) = 61 / (60 + r) +score = Σ_arm weight_arm × rank_weight(rank_arm) / Σ_arm weight_arm, rank_weight(r) = 61 / (60 + r) ``` -A row ranked first by both arms scores 1.0; a row one arm never retrieved contributes 0 for that arm, so an arm's top hit still scores 0.5 and stays visible next to rows deep in the other arm. The fused ordering is final: no diversity selection runs on the hybrid path. +The vector arm weighs 1.0, the lexical arm `LILBEE_LEXICAL_FUSION_WEIGHT`, the title arm `LILBEE_TITLE_SEARCH_WEIGHT`. A row ranked first by every arm scores 1.0; a row only one arm retrieved scores that arm's share of the total weight, so a peaked single-arm hit stays visible next to rows deep in the other arms. The fused ordering is final: no diversity selection runs on the hybrid path. -- **Why rank fusion and not score fusion**: a convex combination of normalized raw scores (`alpha × vector_similarity + (1 − alpha) × normalized_bm25`) was tried here and regressed graded precision about 20% against RRF, at every blend weight. Cosine similarities sit in a high narrow band, giving every dense neighbor a score floor that crowds out lexically-certain rows; ranks are scale-free, so neither arm's score distribution can drown the other. (On the rank-vs-score question, see Bruch et al. 2024, "[An Analysis of Fusion Functions for Hybrid Retrieval](https://dl.acm.org/doi/10.1145/3596512)".) Arm depth matters as much as the formula: a row one arm is certain about scores a fixed 0.5, while rows both arms rank mid-pool accumulate two contributions, so deep candidate pools flood the fused top-k with both-arm mediocrity; arms therefore stay `top_k` deep. +- **Adaptive fusion** (`LILBEE_ADAPTIVE_FUSION`, on by default): the lexical arm's weight is scaled per query by how peaked the vector ranking is, so a confident dense arm downweights lexical and a flat one keeps it. `LILBEE_LEXICAL_FUSION_WEIGHT` is the ceiling the rule scales down from. +- **Why rank fusion and not score fusion**: a convex combination of normalized raw scores (`alpha × vector_similarity + (1 − alpha) × normalized_bm25`) was tried here and regressed graded precision about 20% against RRF, at every blend weight. Cosine similarities sit in a high narrow band, giving every dense neighbor a score floor that crowds out lexically-certain rows; ranks are scale-free, so neither arm's score distribution can drown the other. (On the rank-vs-score question, see Bruch et al. 2024, "[An Analysis of Fusion Functions for Hybrid Retrieval](https://dl.acm.org/doi/10.1145/3596512)".) Arm depth matters as much as the formula: rows both arms rank mid-pool accumulate two contributions, so deep candidate pools flood the fused top-k with both-arm mediocrity; arms therefore stay `top_k` deep. - **Why no MMR on this path**: for lexical queries the relevant passages are often mutually similar (they quote the same identifiers), which is exactly what diversity selection penalizes; running MMR over the fused pool measurably traded relevant lexical hits for diverse off-topic neighbors. MMR still runs on the vector-only fallback path. - **Canonical score**: every search path sets `SearchChunk.score`, and every downstream stage (sorting, filtering, greedy set cover, concept boost, reranker blending) compares only that field. `distance`, `bm25_score`, and the legacy `relevance_score` remain as provenance. -- **Abstention**: the canonical score is [0, 1] with fixed meaning (0.5 = top of one arm), so `min_relevance_score` is a usable floor: when every retrieved chunk falls below it, ask refuses instead of feeding noise as context. Raw RRF sums (~0.016-0.033 total range) made any threshold meaningless. +- **Abstention**: `min_relevance_score` gates on the fused [0, 1] score. Because adaptive fusion normalizes against a per-query weight total, the score is a within-query ranking signal, not a value comparable across queries, so treat the threshold as a coarse floor: when every retrieved chunk falls below it, ask refuses instead of feeding noise as context. Raw RRF sums (~0.016-0.033 total range) made any threshold meaningless. - **max_distance** applies to rows whose *only* signal is a far vector match; a row the BM25 arm also matched keeps its standing regardless of distance, since dropping it would re-bury exactly the identifier hits fusion exists to preserve. - **When it helps**: queries with specific terms, function names, error messages, exact phrases, document identifiers. diff --git a/docs/usage.md b/docs/usage.md index 71f5bdc99..a62c0b9a3 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -812,8 +812,8 @@ something feels off. | `LILBEE_TITLE_SEARCH_WEIGHT` | `0.5` | Title arm weight in rank fusion (1.0 = equal voice with the vector and text arms) | | `LILBEE_LEXICAL_FUSION_WEIGHT` | `1.0` | BM25 arm weight in rank fusion (1.0 = equal to the vector arm; lower it when a strong embedder should dominate a corpus where the lexical arm adds noise) | | `LILBEE_ADAPTIVE_FUSION` | `true` | Scale the BM25 arm's weight per query by how confident the vector arm is (a peaked dense ranking downweights lexical, a flat one keeps it) instead of using a fixed `LILBEE_LEXICAL_FUSION_WEIGHT`. That weight becomes the ceiling the adaptive rule scales down from. Set to `false` to pin the fixed weight | -| `LILBEE_ADAPTIVE_FUSION_MARGIN` | `0.15` | Vector-similarity margin (top hit minus the field) at which adaptive fusion fully silences the BM25 arm; smaller values downweight lexical more aggressively | -| `LILBEE_FILTER_STRUCTURAL_CHUNKS` | `true` | Drop tables-of-contents and cover/title pages from search results. The title and table arms surface these document-structure chunks, which never answer a question and only dilute context precision. Set to `false` to keep them | +| `LILBEE_ADAPTIVE_FUSION_MARGIN` | `0.15` | Vector-similarity margin (top hit minus the field) at which adaptive fusion fully silences the BM25 arm; smaller values downweight lexical more aggressively. `0` disables adaptation, leaving the lexical arm at its full fixed weight | +| `LILBEE_FILTER_STRUCTURAL_CHUNKS` | `false` | Drop tables-of-contents and classification-banner cover/title pages from search results. Off by default: an evaluation A/B found the filter net-negative, since its cover-page heuristic (which fires only on classified-document banners) also caught short banner-carrying body pages. When on, a query-matched or top-ranked page is never dropped. Re-validate per corpus before turning it on | | `LILBEE_ADAPTIVE_THRESHOLD` | `false` | Widen the distance threshold step by step when too few results pass, on the vector-only fallback path (no effect on the default hybrid search) | | `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 | diff --git a/src/lilbee/app/settings_map.py b/src/lilbee/app/settings_map.py index 47b4486b4..acda7bca4 100644 --- a/src/lilbee/app/settings_map.py +++ b/src/lilbee/app/settings_map.py @@ -878,7 +878,7 @@ def get_default(key: str) -> object: bool, nullable=False, group=SettingGroup.RETRIEVAL, - help_text="Drop tables-of-contents and cover pages from search results", + help_text="Drop tables-of-contents and classification-banner cover pages from results", ), "history_rewrite": SettingDef( bool, diff --git a/src/lilbee/core/config/model.py b/src/lilbee/core/config/model.py index b3068cd88..a514b79f4 100644 --- a/src/lilbee/core/config/model.py +++ b/src/lilbee/core/config/model.py @@ -95,10 +95,11 @@ class Config(BaseSettings): max_embed_chars: int = Field(default=2000, ge=1) top_k: int = ConfigField(default=12, ge=1, writable=True) max_distance: float = ConfigField(default=0.75, ge=0.0, writable=True) - # Abstention floor against the canonical [0, 1] relevance score - # (0.0 = no filtering). When every retrieved chunk falls below it, ask - # refuses instead of feeding noise as context. On the fused reciprocal-rank - # scale an arm's top hit scores 0.5, so useful floors start around 0.4. + # Abstention floor against the [0, 1] fused relevance score (0.0 = no + # filtering). When every retrieved chunk falls below it, ask refuses instead + # of feeding noise as context. The fused score normalizes against a per-query + # weight total under adaptive fusion, so this is a coarse floor, not a value + # that means the same thing on every query; tune it against your own corpus. min_relevance_score: float = ConfigField(default=0.0, ge=0.0, writable=True) adaptive_threshold: bool = ConfigField(default=False, writable=True) rag_system_prompt: str = ConfigField( @@ -210,28 +211,24 @@ class Config(BaseSettings): # the retrieval benchmark, not guessed here. lexical_fusion_weight: float = ConfigField(default=1.0, ge=0.0, le=1.0, writable=True) - # Adaptive fusion: instead of a fixed lexical_fusion_weight, scale the BM25 - # arm per query by how confident the vector arm is (a peaked dense ranking - # downweights lexical, a flat one keeps it). On by default: the retrieval - # benchmark found no single fixed weight wins every corpus, and adaptive at - # margin 0.15 beat the fixed-weight default on all three BEIR sets tested - # (biggest gain on the corpus a fixed lexical arm hurt most). lexical_fusion_ - # weight is the ceiling the adaptive rule scales down from; adaptive_fusion_ - # margin is the vector-similarity margin at which the lexical arm is fully - # silenced (smaller = more aggressive). Set adaptive_fusion=false to pin the - # fixed weight instead. + # Adaptive fusion: scale the BM25 arm per query by vector-arm confidence + # instead of a fixed lexical_fusion_weight (a peaked dense ranking downweights + # lexical, a flat one keeps it). On by default. lexical_fusion_weight is the + # ceiling the rule scales down from. Set adaptive_fusion=false to pin the + # fixed weight. adaptive_fusion: bool = ConfigField(default=True, writable=True) + + # Vector-similarity margin at which the lexical arm is fully silenced; smaller + # = more aggressive downweighting. 0 disables adaptation entirely (the lexical + # arm keeps its full fixed weight). adaptive_fusion_margin: float = ConfigField(default=0.15, ge=0.0, le=2.0, writable=True) - # Drop tables-of-contents and cover/title pages from search results. The - # title and table arms surface these document-structure chunks, which usually - # do not answer a question. OFF by default: a UFO-corpus A/B (bb-lenb) found - # the filter net-negative -- its cover-page heuristic also fires on short, - # classification-banner government body pages, dropping content the answer - # needed and hurting faithfulness. When on, a query-matched or top-ranked - # page is never dropped (searcher.search), so the removal is limited to - # structural chunks the query did not actually hit. Re-validate per corpus - # before turning it on. + # Drop tables-of-contents and classification-banner cover/title pages from + # search results. OFF by default: an evaluation A/B on a government-document + # corpus found the filter net-negative, because its cover-page heuristic also + # fires on short banner-carrying body pages. When on, a query-matched or + # top-ranked page is never dropped (searcher.search), so removal is limited to + # structural chunks the query did not hit. Re-validate per corpus before use. filter_structural_chunks: bool = ConfigField(default=False, writable=True) # Chunk count at/above which sync builds an approximate (ANN) vector index diff --git a/src/lilbee/data/ingest/types.py b/src/lilbee/data/ingest/types.py index 8ea4da8b1..f66999efc 100644 --- a/src/lilbee/data/ingest/types.py +++ b/src/lilbee/data/ingest/types.py @@ -73,8 +73,7 @@ class ChunkRecord(TypedDict): chunk: str chunk_index: int vector: list[float] - # Stamped once per document by the pipeline (see _produce_records) rather - # than set by every producer, so the derivation lives in one place. + # Stamped once per document by the pipeline (see _produce_records). title: NotRequired[str] diff --git a/src/lilbee/data/store/core.py b/src/lilbee/data/store/core.py index 018d04e0e..595ec8ea9 100644 --- a/src/lilbee/data/store/core.py +++ b/src/lilbee/data/store/core.py @@ -108,16 +108,26 @@ def _drop_unsupported_far_rows( _MAX_FILTER_ITERATIONS = 20 # safety cap to prevent runaway loops -def _sanitize_fts_query(text: str) -> str: - """Strip the double quotes that make LanceDB parse an FTS query as a phrase. - - lilbee's FTS indexes carry no token positions (bb-rqr8: positions overflow - LanceDB's list encoding on a large corpus), so a phrase query would error. - Replacing quotes with spaces turns a quoted span into plain terms -- which - is what every lilbee FTS call wants anyway, since none needs exact-phrase - matching. +def _lexical_rows( + table: lancedb.table.Table, + query_text: str, + limit: int, + chunk_type: ChunkType | None, + column: str = "chunk", +) -> list[SearchChunk]: + """BM25 rows for *query_text* over a single FTS *column*. + + ``MatchQuery`` pins the column and matches plain terms, so an unpinned search + cannot widen to the title index and a quoted span cannot reach LanceDB as a + phrase (which the positionless index rejects). This is the one place FTS + queries are built; every arm goes through it. """ - return text.replace('"', " ") + from lancedb.query import MatchQuery + + query = table.search(MatchQuery(query_text, column), query_type="fts").limit(limit) + if chunk_type: + query = query.where(_chunk_type_predicate(chunk_type)) + return [SearchChunk(**r) for r in query.to_list()] # Vector ANN index. IVF_PQ compresses vectors so search scales to millions; @@ -441,11 +451,10 @@ def ensure_fts_index(self) -> None: return try: if _has_fts_index(table): - # The index exists and already serves queries, so mark hybrid - # ready BEFORE the optimize: optimize() only folds new rows and - # compacts, and on a large corpus it can hit a LanceDB encoding - # bug and raise. Letting that failure fall through would leave - # _fts_ready False and silently drop every query to vector-only. + # Mark hybrid ready BEFORE optimize(): the existing index + # already serves queries, and optimize() can raise on a large + # corpus. A failure here must not drop every query to + # vector-only. self._fts_ready = True try: # One optimize folds new rows into every index on the @@ -458,14 +467,10 @@ def ensure_fts_index(self) -> None: exc_info=True, ) else: - # No token positions: with_position=True makes the chunk - # index's trigram position lists overflow LanceDB's list - # encoding on optimize() for a large corpus, which then - # serves BM25 un-optimized all session (bb-rqr8). Positions - # only serve exact-phrase queries, which lilbee never issues - # -- query text is sanitized of the quotes that would make - # LanceDB parse it as a phrase (see _sanitize_fts_query), so - # a positionless index never errors on one. + # Positionless: with_position=True overflows LanceDB's list + # encoding on optimize() for a large corpus. Positions only + # serve exact-phrase queries, which lilbee never issues (FTS + # queries match plain terms), so the index never needs them. table.create_fts_index("chunk", replace=False, with_position=False) self._fts_ready = True log.debug("FTS index created on '%s'", CHUNKS_TABLE) @@ -482,8 +487,7 @@ def _ensure_title_fts_unlocked(self, table: lancedb.table.Table) -> None: if _TITLE_COLUMN not in table.schema.names or _has_fts_index(table, _TITLE_COLUMN): return try: - # Positionless for the same reason as the chunk index (bb-rqr8); - # queries are sanitized of phrase quotes so it never errors. + # Positionless for the same reason as the chunk index. table.create_fts_index(_TITLE_COLUMN, replace=False, with_position=False) log.debug("Title FTS index created on '%s'", CHUNKS_TABLE) except Exception: @@ -591,16 +595,7 @@ def bm25_probe( if not self._fts_ready: return [] try: - # Pin the chunk column: once a title FTS index exists, an unpinned - # FTS query searches every indexed column, which would silently - # widen the probe's semantics. - query = table.search( - _sanitize_fts_query(query_text), query_type="fts", fts_columns="chunk" - ) - if chunk_type: - query = query.where(_chunk_type_predicate(chunk_type)) - rows = query.limit(top_k).to_list() - results = [SearchChunk(**r) for r in rows] + results = _lexical_rows(table, query_text, top_k, chunk_type) norms = normalized_bm25([r.bm25_score or 0.0 for r in results]) return [ r.model_copy(update={"score": norm}) for r, norm in zip(results, norms, strict=True) @@ -699,18 +694,8 @@ def _fts_arm( limit: int, chunk_type: ChunkType | None, ) -> list[SearchChunk]: - """BM25-arm candidates over the chunk text. - - The column is pinned: once a title FTS index exists, an unpinned FTS - query searches every indexed column, which would double-count title - tokens this arm's sibling already scores. - """ - query = table.search( - _sanitize_fts_query(query_text), query_type="fts", fts_columns="chunk" - ).limit(limit) - if chunk_type: - query = query.where(_chunk_type_predicate(chunk_type)) - return [SearchChunk(**r) for r in query.to_list()] + """BM25-arm candidates over the chunk text.""" + return _lexical_rows(table, query_text, limit, chunk_type) def _title_arm( self, @@ -726,12 +711,7 @@ def _title_arm( """ if not _has_fts_index(table, _TITLE_COLUMN): return [] - query = table.search( - _sanitize_fts_query(query_text), query_type="fts", fts_columns=_TITLE_COLUMN - ).limit(limit) - if chunk_type: - query = query.where(_chunk_type_predicate(chunk_type)) - return [SearchChunk(**r) for r in query.to_list()] + return _lexical_rows(table, query_text, limit, chunk_type, column=_TITLE_COLUMN) def _hybrid_search( self, @@ -742,24 +722,24 @@ def _hybrid_search( max_distance: float, chunk_type: ChunkType | None = None, ) -> list[SearchChunk]: - """Dual-arm retrieval fused by reciprocal rank; the fused ordering is final. + """Multi-arm retrieval fused by weighted reciprocal rank; the fused ordering is final. + + A vector arm and a chunk-BM25 arm always run; a title-BM25 arm joins + when ``cfg.title_search`` is on. Each row's fused score is the + weight-normalized sum of its arm contributions: the vector arm has + weight 1.0, the lexical arm ``cfg.lexical_fusion_weight`` (scaled per + query when ``cfg.adaptive_fusion`` is on), the title arm + ``cfg.title_search_weight``. So a row a single peaked arm is certain + about scores that arm's share of the total weight, not a fixed 0.5. Each arm fetches exactly ``top_k`` rows. Deeper pools measurably hurt - rank fusion: rows a single arm is certain about score a fixed 0.5, - while rows both arms rank mid-pool accumulate two contributions, so - widening the arms floods the fused top-k with both-arm mediocrity and - buries single-arm certainty (lexical identifier hits above all). - - No diversity selection runs here either. For lexical queries the - relevant passages are often mutually similar (they quote the same - identifiers), which is exactly what MMR penalizes: selecting the - fused pool through MMR measurably traded relevant lexical hits for - diverse off-topic neighbors on graded evaluation. - - When ``cfg.title_search`` is on, a third BM25 arm over document - titles joins the fusion at ``cfg.title_search_weight``; title rows - carry ``bm25_score``, so a title match counts as lexical support - for the distance exemption like any other lexical hit. + rank fusion by flooding the fused top-k with both-arm mediocrity and + burying single-arm certainty (lexical identifier hits above all). No + MMR runs here: lexical passages are often mutually similar, which MMR + penalizes, trading relevant hits for off-topic neighbors. + + Title rows carry ``bm25_score``, so a title match counts as lexical + support for the distance exemption like any other lexical hit. """ title_rows: list[SearchChunk] = [] if self._config.title_search: diff --git a/src/lilbee/data/store/fusion.py b/src/lilbee/data/store/fusion.py index 727b6819a..c5195c82f 100644 --- a/src/lilbee/data/store/fusion.py +++ b/src/lilbee/data/store/fusion.py @@ -12,6 +12,12 @@ dominate); the optional title arm weighs ``title_weight``. Weights rescale the shares without leaving the canonical range. +When adaptive fusion varies ``lexical_weight`` per query (or the title arm +toggles on whether it returned rows), the total weight the score normalizes +against changes query to query, so a fused score is a within-query ranking +signal, not a value comparable across queries. ``min_relevance_score`` is +therefore a coarse floor, not a fixed cross-query threshold. + Rank fusion is deliberate. A convex combination of normalized raw scores (``alpha * vector_similarity + (1 - alpha) * normalized_bm25``) was tried here and measurably regressed graded precision: cosine similarities sit @@ -50,14 +56,13 @@ def adaptive_lexical_weight( A *peaked* vector ranking -- a top hit standing well clear of the field -- means the dense embedder already located the answer and the lexical arm - mostly adds term-match noise (the FiQA case in the retrieval benchmark). A - *flat* ranking means dense is unsure and BM25's exact-term matching is worth - trusting (NFCorpus, SciFact). The confidence signal is the margin between - the top similarity and the mean of the rest, divided by *margin_scale*: at or - above that margin the lexical arm is fully silenced, at zero margin it keeps - its full *base_weight*, and it scales linearly between. Returns *base_weight* - unchanged when there is nothing to measure (fewer than two scored rows) or - the feature is disabled (*margin_scale* <= 0). + mostly adds term-match noise. A *flat* ranking means dense is unsure and + BM25's exact-term matching is worth trusting. The confidence signal is the + margin between the top similarity and the mean of the rest, divided by + *margin_scale*: at or above that margin the lexical arm is fully silenced, + at zero margin it keeps its full *base_weight*, and it scales linearly + between. Returns *base_weight* unchanged when there is nothing to measure + (fewer than two scored rows) or when *margin_scale* <= 0 (adaptation off). """ if margin_scale <= 0: return base_weight @@ -138,7 +143,11 @@ def fuse_arms( total_weight = 1.0 + lexical_weight + (title_weight if title_rows else 0.0) merged: dict[tuple[str, int], SearchChunk] = {} _merge_arm(merged, vector_rows, 1.0 / total_weight) - _merge_arm(merged, fts_rows, lexical_weight / total_weight) - if title_rows: + # A zero-weight arm contributes nothing, so skip it rather than folding in + # zero-score rows that would still carry lexical provenance (and its + # downstream distance/structural exemptions) on no real support. + if lexical_weight > 0: + _merge_arm(merged, fts_rows, lexical_weight / total_weight) + if title_rows and title_weight > 0: _merge_arm(merged, title_rows, title_weight / total_weight) return sorted(merged.values(), key=lambda r: r.score or 0.0, reverse=True) diff --git a/src/lilbee/retrieval/query/dedup.py b/src/lilbee/retrieval/query/dedup.py index a1159a400..fac6a0d47 100644 --- a/src/lilbee/retrieval/query/dedup.py +++ b/src/lilbee/retrieval/query/dedup.py @@ -105,11 +105,12 @@ def filter_results( ) -> list[SearchChunk]: """Drop results below min_relevance_score or above max_distance. - ``min_relevance_score`` gates on the canonical [0, 1] score, which is what - makes an abstention threshold possible. ``max_distance`` additionally drops - rows whose only signal is a far vector match (a row with lexical support - keeps its standing regardless of distance). Pass max_distance=0 to disable - distance filtering. + ``min_relevance_score`` gates on the [0, 1] fused score. Under adaptive + fusion that score normalizes against a per-query weight total, so treat the + threshold as a coarse floor rather than a value that means the same thing on + every query. ``max_distance`` additionally drops rows whose only signal is a + far vector match (a row with lexical support keeps its standing regardless of + distance). Pass max_distance=0 to disable distance filtering. """ if max_distance <= 0 and min_relevance_score <= 0: return results diff --git a/src/lilbee/retrieval/query/searcher.py b/src/lilbee/retrieval/query/searcher.py index f16aaa3ee..7c3446a14 100644 --- a/src/lilbee/retrieval/query/searcher.py +++ b/src/lilbee/retrieval/query/searcher.py @@ -624,15 +624,14 @@ def search( # HTTP, and MCP copies of a bare distance cutoff dropped both-arm # rows the fusion layer deliberately keeps past max_distance. results = filter_results(results, self._config.max_distance) - # Drop tables-of-contents and cover pages: the title and table arms - # surface them, but they never answer a question and only dilute - # context precision (bb-pkn6). Filtered from the top_k*2 candidate - # buffer, so enough real passages remain for the downstream trim. + # Drop tables-of-contents and cover pages that only the vector arm + # surfaced: they dilute context precision without answering a question. + # Filtered from the top_k*2 candidate buffer so enough real passages + # remain for the downstream trim. if self._config.filter_structural_chunks: - # Never drop a page the query actually hit: a lexical (BM25) match or - # the top-ranked result is content the answer may need, whatever its - # shape, so only genuinely-unhit structural chunks are removed - # (bb-lenb: the filter was dropping needed body pages). + # A lexical (BM25 or title) hit or the top-ranked row is content the + # answer may need, whatever its shape, so it is never dropped; only + # structural chunks the lexical arms did not support are removed. results = [ r for i, r in enumerate(results) diff --git a/src/lilbee/retrieval/query/structural.py b/src/lilbee/retrieval/query/structural.py index e01e50b47..493f56bd6 100644 --- a/src/lilbee/retrieval/query/structural.py +++ b/src/lilbee/retrieval/query/structural.py @@ -1,13 +1,11 @@ """Detect document-structure chunks that dilute retrieval precision. -Tables of contents, cover/title pages, and bare reference lists carry a -document's title and section words, so the title arm and table extraction -surface them, but they never *answer* a substantive question. The UFO A/B -(bb-5bjp / bb-pkn6) traced a real context-precision drop to exactly these -chunks entering the retrieved set. The detector is deliberately conservative: -it protects recall/faithfulness by only flagging chunks that are unambiguously -structural, so a false negative (some noise slips through) is preferred to a -false positive (dropping real content). +Flags two classes: tables of contents (generic), and classification-banner +cover/title pages. Both carry a document's title and section words but never +*answer* a substantive question. The detector is deliberately conservative: +it only flags chunks that are unambiguously structural, so a false negative +(some noise slips through) is preferred to a false positive (dropping real +content). """ from __future__ import annotations @@ -26,10 +24,9 @@ _TOC_RATIO = 0.30 # Cover/title-page gates: a title page is very short with essentially no prose. -# These are deliberately tight -- a UFO-corpus A/B (bb-lenb) showed looser gates -# firing on short, classification-banner government BODY pages and dropping the -# content the answer needed, so a real body page's word count or its first full -# sentence must take it out of scope. +# Deliberately tight -- looser gates fire on short banner-carrying body pages and +# drop content the answer needs, so a real body page's word count or its first +# full sentence must take it out of scope. _COVER_MAX_WORDS = 60 _COVER_MAX_SENTENCES = 1 _COVER_CAPS_RATIO = 0.30 diff --git a/tests/test_query.py b/tests/test_query.py index 24bbb9a25..4eb71b33c 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -505,7 +505,7 @@ def test_far_vector_only_rows_filtered_at_search(self, mock_svc): def test_structural_chunk_the_query_missed_is_dropped(self, mock_svc): """With the filter on, a lower-ranked TOC the query did not hit is - dropped, while the real answer passage survives (bb-pkn6).""" + dropped, while the real answer passage survives.""" cfg.filter_structural_chunks = True real = _make_result(source="body.pdf", distance=0.1, chunk="Real answer prose.") toc = _make_result(source="toc.pdf", distance=0.4, chunk=self._TOC_CHUNK) @@ -515,7 +515,7 @@ def test_structural_chunk_the_query_missed_is_dropped(self, mock_svc): def test_structural_filter_keeps_query_matched_page(self, mock_svc): """A page the query lexically hit is never dropped, whatever its shape: - it is content the answer may need (bb-lenb).""" + it is content the answer may need.""" cfg.filter_structural_chunks = True real = _make_result(source="body.pdf", distance=0.1, chunk="Real answer prose.") hit = _make_result(source="toc.pdf", distance=0.4, bm25_score=9.0, chunk=self._TOC_CHUNK) @@ -523,14 +523,14 @@ def test_structural_filter_keeps_query_matched_page(self, mock_svc): assert "toc.pdf" in [r.source for r in get_services().searcher.search("q")] def test_structural_filter_keeps_top_hit(self, mock_svc): - """The top-ranked result is never dropped, whatever its shape (bb-lenb).""" + """The top-ranked result is never dropped, whatever its shape.""" cfg.filter_structural_chunks = True toc = _make_result(source="toc.pdf", distance=0.05, chunk=self._TOC_CHUNK) mock_svc.store.search.return_value = [toc] assert [r.source for r in get_services().searcher.search("q")] == ["toc.pdf"] def test_structural_filter_off_by_default_keeps_toc(self, mock_svc): - """Off by default (bb-lenb: net-negative on the UFO corpus).""" + """Off by default: the A/B found it net-negative on the eval corpus.""" assert cfg.filter_structural_chunks is False toc = _make_result(source="toc.pdf", distance=0.4, chunk=self._TOC_CHUNK) real = _make_result(source="body.pdf", distance=0.1, chunk="Real answer prose.") diff --git a/tests/test_store.py b/tests/test_store.py index ffad94830..a54e2204d 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -193,22 +193,16 @@ def test_first_call_creates_without_replace(self, store): # Verify replace was NOT True (would defeat the purpose of incremental) assert all(c.kwargs.get("replace") is False for c in create_spy.call_args_list) # Both indexes are positionless: with_position=True overflows LanceDB's - # list encoding on a large corpus (bb-rqr8), and no lilbee query needs - # exact-phrase matching. + # list encoding on a large corpus, and no lilbee query needs exact-phrase + # matching. assert all(c.kwargs.get("with_position") is False for c in create_spy.call_args_list) - def test_sanitize_fts_query_replaces_quotes_with_spaces(self): - from lilbee.data.store.core import _sanitize_fts_query - - assert _sanitize_fts_query('"exact phrase" and terms') == " exact phrase and terms" - assert _sanitize_fts_query("plain terms") == "plain terms" - - def test_fts_quoted_query_is_sanitized_not_a_phrase(self, store): + def test_fts_quoted_query_matches_terms_not_a_phrase(self, store): """A quoted query must return term matches, not raise on the positionless index. - The chunk index carries no token positions (bb-rqr8), so a phrase query - would error. bm25_probe strips the quotes (_sanitize_fts_query), turning - the quoted span into plain terms that still match. + The chunk index carries no token positions, so a phrase query would + error. FTS goes through ``MatchQuery``, which matches the quoted span's + plain terms instead of parsing it as a phrase. """ store.add_chunks(_make_records()) store.ensure_fts_index() @@ -695,26 +689,50 @@ def test_hybrid_search_with_fts_index(self, store, test_config): scores = [r.score for r in results] assert all(0.0 <= s <= 1.0 for s in scores) - def test_adaptive_fusion_gates_the_hybrid_search(self, store, test_config): - """With adaptive_fusion on (the default), hybrid search derives the - lexical weight per query from the vector arm instead of the fixed value.""" + def test_adaptive_fusion_feeds_a_derived_weight_to_fusion(self, store, test_config): + """With adaptive_fusion on (the default), the per-query weight from + adaptive_lexical_weight -- fed the configured margin -- is what reaches + fuse_arms, not the fixed config value. Deleting the adaptive branch would + fail this, unlike a smoke test that only checks the score range.""" assert test_config.adaptive_fusion is True # shipped default + test_config.adaptive_fusion_margin = 0.42 store.add_chunks(_make_records()) store.ensure_fts_index() query_vec = [0.5] * test_config.embedding_dim - results = store.search(query_vec, top_k=3, query_text="chunk number") + from lilbee.data.store import core as store_core + + with ( + mock.patch.object(store_core, "adaptive_lexical_weight", return_value=0.123) as adapt, + mock.patch.object(store_core, "fuse_arms", wraps=store_core.fuse_arms) as fuse, + ): + results = store.search(query_vec, top_k=3, query_text="chunk number") + adapt.assert_called_once() + # adaptive_lexical_weight(vector_rows, base_weight, margin): the base is + # lexical_fusion_weight and the margin is the configured value. + assert adapt.call_args.args[1] == pytest.approx(test_config.lexical_fusion_weight) + assert adapt.call_args.args[2] == pytest.approx(0.42) + assert fuse.call_args.kwargs["lexical_weight"] == pytest.approx(0.123) assert len(results) > 0 - assert all(r.score is not None and 0.0 <= r.score <= 1.0 for r in results) - def test_fixed_fusion_when_adaptive_disabled(self, store, test_config): - """Opting out of adaptive fusion pins the fixed lexical_fusion_weight.""" + def test_fixed_fusion_pins_the_config_weight(self, store, test_config): + """Opting out of adaptive fusion skips adaptive_lexical_weight and pins + the fixed lexical_fusion_weight into fuse_arms; a non-default weight must + reach fusion verbatim.""" test_config.adaptive_fusion = False + test_config.lexical_fusion_weight = 0.3 store.add_chunks(_make_records()) store.ensure_fts_index() query_vec = [0.5] * test_config.embedding_dim - results = store.search(query_vec, top_k=3, query_text="chunk number") + from lilbee.data.store import core as store_core + + with ( + mock.patch.object(store_core, "adaptive_lexical_weight") as adapt, + mock.patch.object(store_core, "fuse_arms", wraps=store_core.fuse_arms) as fuse, + ): + results = store.search(query_vec, top_k=3, query_text="chunk number") + adapt.assert_not_called() + assert fuse.call_args.kwargs["lexical_weight"] == pytest.approx(0.3) assert len(results) > 0 - assert all(r.score is not None and 0.0 <= r.score <= 1.0 for r in results) def test_fallback_to_vector_when_no_query_text(self, store, test_config): records = _make_records() @@ -2424,6 +2442,20 @@ def test_title_arm_surfaces_title_only_match(self, store, test_config): assert matched assert all(r.bm25_score is not None for r in matched) + def test_title_search_weight_reaches_fusion(self, store, test_config): + """A non-default title_search_weight is threaded into fuse_arms, not + hardcoded: the config value is what weights the title arm.""" + test_config.title_search = True + test_config.title_search_weight = 0.2 + store.add_chunks(_titled_records("a.pdf", 2, title="zebra manifesto")) + store.ensure_fts_index() + query_vec = [0.1] * test_config.embedding_dim + from lilbee.data.store import core as store_core + + with mock.patch.object(store_core, "fuse_arms", wraps=store_core.fuse_arms) as fuse: + store.search(query_vec, top_k=4, max_distance=0, query_text="zebra") + assert fuse.call_args.kwargs["title_weight"] == pytest.approx(0.2) + def test_title_arm_off_by_default(self, store, test_config): """With title_search off, a title-only term earns no lexical support.""" assert test_config.title_search is False diff --git a/tests/test_store_fusion.py b/tests/test_store_fusion.py index 15762bf5b..922e923c5 100644 --- a/tests/test_store_fusion.py +++ b/tests/test_store_fusion.py @@ -72,6 +72,37 @@ def test_lexical_only_row_survives_with_score(self): lexical_row = next(r for r in fused if r.source == "catalog_482.pdf") assert lexical_row.score == pytest.approx(0.5) + def test_top_lexical_hit_outranks_all_but_the_top_dense_neighbor(self): + """The pinpoint-document failure mode: an FTS-arm top hit unseen by + the vector arm must rank above every vector row except at most the + vector arm's own number one.""" + vec = [_chunk("noise.md", i, distance=0.1) for i in range(30)] + lex = [_chunk("target.pdf", 0, bm25=30.0), _chunk("noise.md", 0, bm25=3.0)] + fused = fuse_arms(vec, lex) + rank = next(i for i, r in enumerate(fused) if r.source == "target.pdf") + assert rank <= 1 + + def test_rank_order_ignores_score_magnitude(self): + """A wildly stronger BM25 score at rank 2 stays rank 2: fusion is + scale-free by construction.""" + lex = [_chunk("first.md", 0, bm25=5.0), _chunk("second.md", 0, bm25=500.0)] + fused = fuse_arms([], lex) + assert [r.source for r in fused] == ["first.md", "second.md"] + + def test_dedup_keeps_both_provenance_fields(self): + fused = fuse_arms([_chunk("a.md", 0, distance=0.3)], [_chunk("a.md", 0, bm25=9.0)]) + assert len(fused) == 1 + assert fused[0].distance == pytest.approx(0.3) + assert fused[0].bm25_score == pytest.approx(9.0) + + def test_sorted_descending_by_score(self): + fused = fuse_arms( + [_chunk("near.md", 0, distance=0.1), _chunk("far.md", 0, distance=1.5)], + [], + ) + assert [r.source for r in fused] == ["near.md", "far.md"] + assert all(r.score is not None for r in fused) + class TestAdaptiveLexicalWeight: """Per-query lexical weight gated by the vector arm's confidence.""" @@ -134,43 +165,18 @@ def test_default_weight_is_the_historical_equal_voice(self): explicit = {r.source: r.score for r in fuse_arms(vec, lex, lexical_weight=1.0)} assert default == explicit - def test_zero_weight_silences_the_lexical_arm(self): - assert self._lex_row(0.0).score == pytest.approx(0.0) + def test_zero_weight_drops_the_lexical_only_arm(self): + """A fully-silenced lexical arm contributes no rows at all, so a + BM25-only hit never enters the pool carrying lexical provenance (and + with it the downstream distance/structural exemptions).""" + vec = [_chunk("noise.md", 0, distance=0.5)] + lex = [_chunk("cat.pdf", 0, bm25=35.0)] + fused = fuse_arms(vec, lex, lexical_weight=0.0) + assert "cat.pdf" not in [r.source for r in fused] def test_lower_weight_shrinks_the_lexical_contribution(self): assert self._lex_row(0.5).score < self._lex_row(1.0).score - def test_top_lexical_hit_outranks_all_but_the_top_dense_neighbor(self): - """The pinpoint-document failure mode: an FTS-arm top hit unseen by - the vector arm must rank above every vector row except at most the - vector arm's own number one.""" - vec = [_chunk("noise.md", i, distance=0.1) for i in range(30)] - lex = [_chunk("target.pdf", 0, bm25=30.0), _chunk("noise.md", 0, bm25=3.0)] - fused = fuse_arms(vec, lex) - rank = next(i for i, r in enumerate(fused) if r.source == "target.pdf") - assert rank <= 1 - - def test_rank_order_ignores_score_magnitude(self): - """A wildly stronger BM25 score at rank 2 stays rank 2: fusion is - scale-free by construction.""" - lex = [_chunk("first.md", 0, bm25=5.0), _chunk("second.md", 0, bm25=500.0)] - fused = fuse_arms([], lex) - assert [r.source for r in fused] == ["first.md", "second.md"] - - def test_dedup_keeps_both_provenance_fields(self): - fused = fuse_arms([_chunk("a.md", 0, distance=0.3)], [_chunk("a.md", 0, bm25=9.0)]) - assert len(fused) == 1 - assert fused[0].distance == pytest.approx(0.3) - assert fused[0].bm25_score == pytest.approx(9.0) - - def test_sorted_descending_by_score(self): - fused = fuse_arms( - [_chunk("near.md", 0, distance=0.1), _chunk("far.md", 0, distance=1.5)], - [], - ) - assert [r.source for r in fused] == ["near.md", "far.md"] - assert all(r.score is not None for r in fused) - class TestRegressionMechanism: """Synthetic reproduction of the graded-A/B regression shape: a lexical diff --git a/tests/test_structural.py b/tests/test_structural.py index 4d36d02b1..f3d8fb312 100644 --- a/tests/test_structural.py +++ b/tests/test_structural.py @@ -2,41 +2,41 @@ from lilbee.retrieval.query.structural import is_structural_chunk -# A real table of contents, as seen diluting precision in the UFO A/B (bb-pkn6). +# A table of contents: section titles followed by dot leaders and page numbers. TOC = """Contents A. Executive Summary ................................................. 1 B. Introduction ..................................................... 3 -C. Geographic Trends ................................................ 9 -D. Notable Trends Regarding Propulsion .............................. 10 -E. Flight Safety Issues ............................................. 11 -F. UAS Observations Reported ........................................ 13 +C. Regional Overview ............................................... 9 +D. Budget and Staffing ............................................. 10 +E. Program Findings ................................................ 11 +F. Recommendations ................................................. 13 """ -# A real DoD cover/title page. +# A classification-banner cover/title page. COVER = """UNCLASSIFIED -THE DEPARTMENT OF DEFENSE -ALL-DOMAIN ANOMALY RESOLUTION OFFICE +NATIONAL PROGRAM REVIEW BOARD +OFFICE OF STRATEGIC ASSESSMENT Fiscal Year 2024 Consolidated Annual Report on -Unidentified Anomalous Phenomena +Program Performance and Oversight Information Cut Off Date: 1 JUNE 2024 UNCLASSIFIED """ # Real prose that must NOT be flagged. PROSE = ( - "During the reporting period, AARO received no reports indicating UAP sightings " - "have been associated with any adverse health effects. However, many reports from " - "military witnesses described transient effects. The office continues to evaluate " - "each case against a standardized methodology, and it has resolved the majority of " - "reported incidents as ordinary objects or sensor artifacts." + "During the reporting period, the office received no reports indicating that the " + "program had an adverse effect on regional operations. However, many reports from " + "field staff described transient delays. The office continues to evaluate each case " + "against a standardized methodology, and it has resolved the majority of reported " + "incidents as routine administrative issues." ) # Prose that happens to reference a page and carry a classification header. PROSE_WITH_HEADER = ( - "UNCLASSIFIED. As detailed on page 12, the assessment concluded that the object was " - "a commercial aircraft. The radar track and the electro-optical imagery were " - "consistent, and the case was closed. No anomalous performance was observed at any " - "point during the encounter, which lasted several minutes." + "UNCLASSIFIED. As detailed on page 12, the assessment concluded that the shortfall was " + "a routine scheduling gap. The budget records and the staffing figures were " + "consistent, and the case was closed. No irregularity was observed at any point " + "during the review, which lasted several weeks." ) @@ -63,16 +63,16 @@ def test_single_dot_leader_line_is_not_a_toc(self): def test_short_all_caps_without_classification_is_not_a_cover(self): # A shouting heading with no classification banner is left alone. - assert is_structural_chunk("NOTABLE TRENDS REGARDING PROPULSION AND FLIGHT") is False + assert is_structural_chunk("REGIONAL OVERVIEW AND PROGRAM FINDINGS") is False def test_short_classified_body_page_is_not_a_cover(self): - # A short government body page carries a classification banner and caps - # but real content; its full sentences must keep it out of scope so the - # answer does not lose the page it needs (bb-lenb, the rag-1 failure). + # A short body page carries a classification banner and caps but real + # content; its full sentences must keep it out of scope so the answer + # does not lose the page it needs. body = ( - "UNCLASSIFIED. The AARO assessment concluded the object was a " - "commercial aircraft, and the case was resolved. RADAR and EO/IR " - "data were CONSISTENT across the entire track." + "UNCLASSIFIED. The office assessment concluded the shortfall was a " + "routine scheduling gap, and the case was resolved. BUDGET and STAFFING " + "data were CONSISTENT across the entire review." ) assert is_structural_chunk(body) is False From 0c8597a9c51a15caacaebb2c1a5ebde94aeec26f Mon Sep 17 00:00:00 2001 From: Tobias Perelstein <5562156+tobocop2@users.noreply.github.com> Date: Sun, 19 Jul 2026 20:01:01 -0400 Subject: [PATCH 14/28] Normalize fused scores against a constant weight budget Adaptive fusion computed each search's normalization denominator from that query's own adapted lexical weight and whether its title arm returned rows. Since Searcher issues a separate hybrid search per expansion variant and merges them into one sorted pool, rows from a query whose lexical arm was silenced (denominator 1.0) outscored identical-strength rows from a query whose lexical arm was active (denominator 2.0), independent of relevance. The same drift moved min_relevance_score's meaning query to query. fuse_arms now takes an explicit weight_total, and hybrid search passes the configured weight budget (1 + lexical_fusion_weight + title weight when the title arm is enabled) rather than the per-query outcome. The denominator is a uniform divisor, so ranking within a single search is unchanged; only the cross-search scale is fixed, which is the point. min_relevance_score is a stable floor again. Docstrings, config comments, and architecture.md updated to match; direct fuse_arms callers keep the old per-call default. The expansion-path merge ordering changes as a result (it stops favoring whichever sub-search silenced its lexical arm); the primary single-search ranking is unchanged. --- docs/architecture.md | 4 +-- src/lilbee/core/config/model.py | 6 ++-- src/lilbee/data/store/core.py | 10 ++++++- src/lilbee/data/store/fusion.py | 27 +++++++++++------ src/lilbee/retrieval/query/dedup.py | 11 ++++--- tests/test_store.py | 10 +++++++ tests/test_store_fusion.py | 46 +++++++++++++++++++++++++++++ 7 files changed, 93 insertions(+), 21 deletions(-) diff --git a/docs/architecture.md b/docs/architecture.md index 4ddd3ac9e..efe0934d7 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -686,11 +686,11 @@ score = Σ_arm weight_arm × rank_weight(rank_arm) / Σ_arm weight_arm, rank_we The vector arm weighs 1.0, the lexical arm `LILBEE_LEXICAL_FUSION_WEIGHT`, the title arm `LILBEE_TITLE_SEARCH_WEIGHT`. A row ranked first by every arm scores 1.0; a row only one arm retrieved scores that arm's share of the total weight, so a peaked single-arm hit stays visible next to rows deep in the other arms. The fused ordering is final: no diversity selection runs on the hybrid path. -- **Adaptive fusion** (`LILBEE_ADAPTIVE_FUSION`, on by default): the lexical arm's weight is scaled per query by how peaked the vector ranking is, so a confident dense arm downweights lexical and a flat one keeps it. `LILBEE_LEXICAL_FUSION_WEIGHT` is the ceiling the rule scales down from. +- **Adaptive fusion** (`LILBEE_ADAPTIVE_FUSION`, on by default): the lexical arm's weight is scaled per query by how peaked the vector ranking is, so a confident dense arm downweights lexical and a flat one keeps it. `LILBEE_LEXICAL_FUSION_WEIGHT` is the ceiling the rule scales down from. Scores still normalize against that configured budget (a constant), not the per-query adapted weight, so scores from the separate sub-searches expansion merges stay on one comparable scale. - **Why rank fusion and not score fusion**: a convex combination of normalized raw scores (`alpha × vector_similarity + (1 − alpha) × normalized_bm25`) was tried here and regressed graded precision about 20% against RRF, at every blend weight. Cosine similarities sit in a high narrow band, giving every dense neighbor a score floor that crowds out lexically-certain rows; ranks are scale-free, so neither arm's score distribution can drown the other. (On the rank-vs-score question, see Bruch et al. 2024, "[An Analysis of Fusion Functions for Hybrid Retrieval](https://dl.acm.org/doi/10.1145/3596512)".) Arm depth matters as much as the formula: rows both arms rank mid-pool accumulate two contributions, so deep candidate pools flood the fused top-k with both-arm mediocrity; arms therefore stay `top_k` deep. - **Why no MMR on this path**: for lexical queries the relevant passages are often mutually similar (they quote the same identifiers), which is exactly what diversity selection penalizes; running MMR over the fused pool measurably traded relevant lexical hits for diverse off-topic neighbors. MMR still runs on the vector-only fallback path. - **Canonical score**: every search path sets `SearchChunk.score`, and every downstream stage (sorting, filtering, greedy set cover, concept boost, reranker blending) compares only that field. `distance`, `bm25_score`, and the legacy `relevance_score` remain as provenance. -- **Abstention**: `min_relevance_score` gates on the fused [0, 1] score. Because adaptive fusion normalizes against a per-query weight total, the score is a within-query ranking signal, not a value comparable across queries, so treat the threshold as a coarse floor: when every retrieved chunk falls below it, ask refuses instead of feeding noise as context. Raw RRF sums (~0.016-0.033 total range) made any threshold meaningless. +- **Abstention**: `min_relevance_score` gates on the fused [0, 1] score. The score normalizes against the configured weight budget (a constant), so an arm's top hit scores a stable share of it and the threshold is a usable floor: when every retrieved chunk falls below it, ask refuses instead of feeding noise as context. Raw RRF sums (~0.016-0.033 total range) made any threshold meaningless. - **max_distance** applies to rows whose *only* signal is a far vector match; a row the BM25 arm also matched keeps its standing regardless of distance, since dropping it would re-bury exactly the identifier hits fusion exists to preserve. - **When it helps**: queries with specific terms, function names, error messages, exact phrases, document identifiers. diff --git a/src/lilbee/core/config/model.py b/src/lilbee/core/config/model.py index 3f293255f..a39f8bd21 100644 --- a/src/lilbee/core/config/model.py +++ b/src/lilbee/core/config/model.py @@ -97,9 +97,9 @@ class Config(BaseSettings): max_distance: float = ConfigField(default=0.75, ge=0.0, writable=True) # Abstention floor against the [0, 1] fused relevance score (0.0 = no # filtering). When every retrieved chunk falls below it, ask refuses instead - # of feeding noise as context. The fused score normalizes against a per-query - # weight total under adaptive fusion, so this is a coarse floor, not a value - # that means the same thing on every query; tune it against your own corpus. + # of feeding noise as context. The fused score normalizes against the + # configured weight budget (a constant), so an arm's top hit scores a stable + # share of it; useful floors start around 0.4. Tune against your own corpus. min_relevance_score: float = ConfigField(default=0.0, ge=0.0, writable=True) adaptive_threshold: bool = ConfigField(default=False, writable=True) rag_system_prompt: str = ConfigField( diff --git a/src/lilbee/data/store/core.py b/src/lilbee/data/store/core.py index 595ec8ea9..eda5eca1d 100644 --- a/src/lilbee/data/store/core.py +++ b/src/lilbee/data/store/core.py @@ -745,18 +745,26 @@ def _hybrid_search( if self._config.title_search: title_rows = self._title_arm(table, query_text, top_k, chunk_type) vector_rows = self._vector_arm(table, query_vector, top_k, chunk_type) - lexical_weight = self._config.lexical_fusion_weight + base_lexical_weight = self._config.lexical_fusion_weight + lexical_weight = base_lexical_weight if self._config.adaptive_fusion: # Gate the lexical arm per query by how peaked the vector ranking is. lexical_weight = adaptive_lexical_weight( vector_rows, lexical_weight, self._config.adaptive_fusion_margin ) + # Normalize against the configured weight budget, not this query's adapted + # weight or whether its title arm happened to return rows, so scores stay + # comparable across the sub-searches Searcher merges (query + variants). + weight_total = 1.0 + base_lexical_weight + ( + self._config.title_search_weight if self._config.title_search else 0.0 + ) fused = fuse_arms( vector_rows, self._fts_arm(table, query_text, top_k, chunk_type), title_rows, lexical_weight=lexical_weight, title_weight=self._config.title_search_weight, + weight_total=weight_total, ) fused = _drop_unsupported_far_rows(fused, max_distance) return fused[:top_k] diff --git a/src/lilbee/data/store/fusion.py b/src/lilbee/data/store/fusion.py index c5195c82f..66952a86d 100644 --- a/src/lilbee/data/store/fusion.py +++ b/src/lilbee/data/store/fusion.py @@ -12,11 +12,11 @@ dominate); the optional title arm weighs ``title_weight``. Weights rescale the shares without leaving the canonical range. -When adaptive fusion varies ``lexical_weight`` per query (or the title arm -toggles on whether it returned rows), the total weight the score normalizes -against changes query to query, so a fused score is a within-query ranking -signal, not a value comparable across queries. ``min_relevance_score`` is -therefore a coarse floor, not a fixed cross-query threshold. +Scores normalize against the configured weight budget, not the per-query +adapted ``lexical_weight`` or whether a given query's title arm returned rows +(see ``fuse_arms``'s *weight_total*). That fixed denominator keeps scores on +one scale across queries and across the sub-searches a caller merges, so +``min_relevance_score`` stays a usable floor. Rank fusion is deliberate. A convex combination of normalized raw scores (``alpha * vector_similarity + (1 - alpha) * normalized_bm25``) was tried @@ -130,6 +130,7 @@ def fuse_arms( *, lexical_weight: float = 1.0, title_weight: float = 1.0, + weight_total: float | None = None, ) -> list[SearchChunk]: """Merge the arms into one list scored by reciprocal rank. @@ -139,15 +140,23 @@ def fuse_arms( arms carry every provenance field (``distance`` from the vector arm, ``bm25_score`` from the FTS arms). The result is sorted by ``score`` descending and deduplicated on ``(source, chunk_index)``. + + *weight_total* is the constant this call normalizes scores against. Pass the + configured weight budget so scores from separate sub-searches -- whose + adaptive *lexical_weight* and per-query title-arm presence differ -- stay on + one comparable scale when a caller merges them. It is a uniform divisor, so + the ranking within this call is unchanged. When omitted it defaults to the + arms present in this call (the standalone behavior). """ - total_weight = 1.0 + lexical_weight + (title_weight if title_rows else 0.0) + if weight_total is None: + weight_total = 1.0 + lexical_weight + (title_weight if title_rows else 0.0) merged: dict[tuple[str, int], SearchChunk] = {} - _merge_arm(merged, vector_rows, 1.0 / total_weight) + _merge_arm(merged, vector_rows, 1.0 / weight_total) # A zero-weight arm contributes nothing, so skip it rather than folding in # zero-score rows that would still carry lexical provenance (and its # downstream distance/structural exemptions) on no real support. if lexical_weight > 0: - _merge_arm(merged, fts_rows, lexical_weight / total_weight) + _merge_arm(merged, fts_rows, lexical_weight / weight_total) if title_rows and title_weight > 0: - _merge_arm(merged, title_rows, title_weight / total_weight) + _merge_arm(merged, title_rows, title_weight / weight_total) return sorted(merged.values(), key=lambda r: r.score or 0.0, reverse=True) diff --git a/src/lilbee/retrieval/query/dedup.py b/src/lilbee/retrieval/query/dedup.py index fac6a0d47..9f5314d99 100644 --- a/src/lilbee/retrieval/query/dedup.py +++ b/src/lilbee/retrieval/query/dedup.py @@ -105,12 +105,11 @@ def filter_results( ) -> list[SearchChunk]: """Drop results below min_relevance_score or above max_distance. - ``min_relevance_score`` gates on the [0, 1] fused score. Under adaptive - fusion that score normalizes against a per-query weight total, so treat the - threshold as a coarse floor rather than a value that means the same thing on - every query. ``max_distance`` additionally drops rows whose only signal is a - far vector match (a row with lexical support keeps its standing regardless of - distance). Pass max_distance=0 to disable distance filtering. + ``min_relevance_score`` gates on the [0, 1] fused score, which normalizes + against the configured weight budget (a constant), so the threshold means the + same thing across queries. ``max_distance`` additionally drops rows whose only + signal is a far vector match (a row with lexical support keeps its standing + regardless of distance). Pass max_distance=0 to disable distance filtering. """ if max_distance <= 0 and min_relevance_score <= 0: return results diff --git a/tests/test_store.py b/tests/test_store.py index a54e2204d..4ee7b1de7 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -712,6 +712,11 @@ def test_adaptive_fusion_feeds_a_derived_weight_to_fusion(self, store, test_conf assert adapt.call_args.args[1] == pytest.approx(test_config.lexical_fusion_weight) assert adapt.call_args.args[2] == pytest.approx(0.42) assert fuse.call_args.kwargs["lexical_weight"] == pytest.approx(0.123) + # The normalization denominator is the configured budget (constant), not + # the adapted 0.123, so scores stay comparable across sub-searches. + assert fuse.call_args.kwargs["weight_total"] == pytest.approx( + 1.0 + test_config.lexical_fusion_weight + ) assert len(results) > 0 def test_fixed_fusion_pins_the_config_weight(self, store, test_config): @@ -2455,6 +2460,11 @@ def test_title_search_weight_reaches_fusion(self, store, test_config): with mock.patch.object(store_core, "fuse_arms", wraps=store_core.fuse_arms) as fuse: store.search(query_vec, top_k=4, max_distance=0, query_text="zebra") assert fuse.call_args.kwargs["title_weight"] == pytest.approx(0.2) + # With the title arm enabled, its weight is part of the constant + # denominator whether or not this query's title arm returned rows. + assert fuse.call_args.kwargs["weight_total"] == pytest.approx( + 1.0 + test_config.lexical_fusion_weight + 0.2 + ) def test_title_arm_off_by_default(self, store, test_config): """With title_search off, a title-only term earns no lexical support.""" diff --git a/tests/test_store_fusion.py b/tests/test_store_fusion.py index 922e923c5..7721378c7 100644 --- a/tests/test_store_fusion.py +++ b/tests/test_store_fusion.py @@ -104,6 +104,52 @@ def test_sorted_descending_by_score(self): assert all(r.score is not None for r in fused) +class TestWeightTotalNormalization: + """A constant reference denominator keeps scores comparable across the + separate sub-searches that Searcher merges. Without it, each sub-search + normalizes by its own per-query total_weight, so a peaked sub-search + (lexical silenced) inflates its rows against a flat one's.""" + + def test_weight_total_pins_the_denominator(self): + # Lexical silenced (weight 0): the vector-only top hit must still be + # scored against the supplied constant denominator, not 1.0. + fused = fuse_arms( + [_chunk("a.md", 0, distance=0.3)], [], lexical_weight=0.0, weight_total=2.0 + ) + assert fused[0].score == pytest.approx(0.5) + + def test_cross_subsearch_scores_share_one_scale(self): + # Same identical-strength vector-only top hit, two sub-searches whose + # adaptive lexical weight differs: with one shared weight_total they land + # on the same score instead of 1.0 vs 0.5. + peaked = fuse_arms( + [_chunk("a.md", 0, distance=0.3)], [], lexical_weight=0.0, weight_total=2.0 + ) + flat = fuse_arms( + [_chunk("a.md", 0, distance=0.3)], [], lexical_weight=1.0, weight_total=2.0 + ) + assert peaked[0].score == pytest.approx(flat[0].score) + + def test_weight_total_preserves_within_call_order(self): + # A uniform denominator is a monotonic rescale, so the ranking inside one + # call is identical to the default per-call normalization. + vec = [_chunk("near.md", 0, distance=0.1), _chunk("far.md", 1, distance=0.9)] + lex = [_chunk("far.md", 1, bm25=30.0)] + default = [r.source for r in fuse_arms(vec, lex, lexical_weight=1.0)] + pinned = [r.source for r in fuse_arms(vec, lex, lexical_weight=1.0, weight_total=2.0)] + assert default == pinned + + def test_weight_total_defaults_to_per_call_when_absent(self): + # Direct callers that omit weight_total keep the original behavior. + no_arg = fuse_arms([_chunk("a.md", 0, distance=0.3)], [_chunk("a.md", 0, bm25=9.0)]) + explicit = fuse_arms( + [_chunk("a.md", 0, distance=0.3)], + [_chunk("a.md", 0, bm25=9.0)], + weight_total=2.0, + ) + assert no_arg[0].score == pytest.approx(explicit[0].score) + + class TestAdaptiveLexicalWeight: """Per-query lexical weight gated by the vector arm's confidence.""" From 6547b54b6d4b0993053a588ef83f693db9850a47 Mon Sep 17 00:00:00 2001 From: Tobias Perelstein <5562156+tobocop2@users.noreply.github.com> Date: Sun, 19 Jul 2026 20:20:32 -0400 Subject: [PATCH 15/28] Fix title-arm isolation, gating, adaptive parity, and upgrade visibility Addresses five PR #557 audit findings on the title/FTS surface: - Title-arm query failure no longer collapses the whole hybrid search to vector-only: _title_arm catches and returns [] like bm25_probe, so a broken title index degrades to no-titles instead of taking down the chunk arm. - The title FTS index is built only when title_search is enabled, instead of mutating every store with a second index nothing queries. - Adaptive fusion now scales the title arm by the same confidence factor as the chunk-BM25 arm; leaving it at full weight re-admitted the lexical signal adaptive fusion had just silenced. adaptive_weight_scale is factored out and applied to both lexical arms. - Migrating a pre-title store logs that documents indexed before the upgrade stay title-blind until a rebuild, with a matching docs note; nothing warned before. - A regression test asserts the neighbor-widened prompt stays under the provider's prompt_token_budget ceiling (the merge already routes widening through the same budget; this pins it). make check pending; targeted suites green. --- docs/usage.md | 2 +- src/lilbee/data/store/core.py | 45 +++++++++++----- src/lilbee/data/store/fusion.py | 31 ++++++----- tests/test_query.py | 36 +++++++++++++ tests/test_store.py | 94 +++++++++++++++++++++++++++------ 5 files changed, 166 insertions(+), 42 deletions(-) diff --git a/docs/usage.md b/docs/usage.md index 9878070b7..6411714dc 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -814,7 +814,7 @@ something feels off. | `LILBEE_RERANK_BLEND` | `true` | Blend reranker scores with retrieval fusion; off = the reranker's own ordering stands | | `LILBEE_HYDE` | `false` | Enable Hypothetical Document Embeddings: an LLM drafts a hypothetical answer, that's embedded, and results are merged with the original query's. Adds ~500 ms per query; helps on vague questions | | `LILBEE_HYDE_WEIGHT` | `0.7` | How much to trust HyDE results relative to the direct query (0.0-1.0) | -| `LILBEE_TITLE_SEARCH` | `false` | Match queries against document titles as a third hybrid-search arm, so a query naming a document by title surfaces its chunks | +| `LILBEE_TITLE_SEARCH` | `false` | Match queries against document titles as a third hybrid-search arm, so a query naming a document by title surfaces its chunks. Documents indexed before this feature existed have no stored title; run `lilbee rebuild` to make them title-searchable | | `LILBEE_TITLE_SEARCH_WEIGHT` | `0.5` | Title arm weight in rank fusion (1.0 = equal voice with the vector and text arms) | | `LILBEE_LEXICAL_FUSION_WEIGHT` | `1.0` | BM25 arm weight in rank fusion (1.0 = equal to the vector arm; lower it when a strong embedder should dominate a corpus where the lexical arm adds noise) | | `LILBEE_ADAPTIVE_FUSION` | `true` | Scale the BM25 arm's weight per query by how confident the vector arm is (a peaked dense ranking downweights lexical, a flat one keeps it) instead of using a fixed `LILBEE_LEXICAL_FUSION_WEIGHT`. That weight becomes the ceiling the adaptive rule scales down from. Set to `false` to pin the fixed weight | diff --git a/src/lilbee/data/store/core.py b/src/lilbee/data/store/core.py index eda5eca1d..40570aad8 100644 --- a/src/lilbee/data/store/core.py +++ b/src/lilbee/data/store/core.py @@ -28,7 +28,7 @@ from lilbee.core.security import validate_path_within from lilbee.runtime.lock import LOCK_TIMEOUT, write_lock -from .fusion import adaptive_lexical_weight, fuse_arms, normalized_bm25, vector_similarity +from .fusion import adaptive_weight_scale, fuse_arms, normalized_bm25, vector_similarity from .lance_helpers import ( _chunk_type_predicate, _has_fts_index, @@ -261,6 +261,13 @@ def _chunks_table(self) -> lancedb.table.Table: table = ensure_table(self.get_db(), CHUNKS_TABLE, self._chunks_schema()) if _TITLE_COLUMN not in table.schema.names: table.add_columns({_TITLE_COLUMN: "CAST(NULL AS STRING)"}) + # Titles are stamped only at ingest, so existing rows stay NULL and + # the title arm cannot see them. Point the user at a rebuild. + log.warning( + "Added the title column to an existing store. Documents indexed " + "before this upgrade have no title, so the title-search arm will " + "not match them until you run `lilbee rebuild`." + ) return table def get_meta(self) -> StoreMeta | None: @@ -474,7 +481,11 @@ def ensure_fts_index(self) -> None: table.create_fts_index("chunk", replace=False, with_position=False) self._fts_ready = True log.debug("FTS index created on '%s'", CHUNKS_TABLE) - self._ensure_title_fts_unlocked(table) + # Only the opt-in title arm needs the title index; without this a + # second FTS index is built (and optimized) on every store even + # though nothing queries it. + if self._config.title_search: + self._ensure_title_fts_unlocked(table) except Exception: log.debug("FTS index ensure failed (empty table?)", exc_info=True) @@ -706,12 +717,18 @@ def _title_arm( ) -> list[SearchChunk]: """BM25-arm candidates over the document title column. - Empty when the store predates the title column or its FTS index: old - indexes keep working, the arm simply contributes nothing there. + Empty when the store predates the title column or its FTS index (old + indexes keep working) and empty on any query-time failure: the optional + title arm must never take down the healthy chunk arm, so its failure + degrades to no-titles, mirroring ``bm25_probe``. """ if not _has_fts_index(table, _TITLE_COLUMN): return [] - return _lexical_rows(table, query_text, limit, chunk_type, column=_TITLE_COLUMN) + try: + return _lexical_rows(table, query_text, limit, chunk_type, column=_TITLE_COLUMN) + except Exception: + log.debug("Title arm search failed; contributing no title rows", exc_info=True) + return [] def _hybrid_search( self, @@ -746,24 +763,28 @@ def _hybrid_search( title_rows = self._title_arm(table, query_text, top_k, chunk_type) vector_rows = self._vector_arm(table, query_vector, top_k, chunk_type) base_lexical_weight = self._config.lexical_fusion_weight + base_title_weight = self._config.title_search_weight lexical_weight = base_lexical_weight + title_weight = base_title_weight if self._config.adaptive_fusion: - # Gate the lexical arm per query by how peaked the vector ranking is. - lexical_weight = adaptive_lexical_weight( - vector_rows, lexical_weight, self._config.adaptive_fusion_margin - ) + # Gate the lexical arms per query by how peaked the vector ranking is. + # The title arm is lexical too, so the same factor scales it; leaving + # it at full weight would re-admit the signal this just silenced. + scale = adaptive_weight_scale(vector_rows, self._config.adaptive_fusion_margin) + lexical_weight = base_lexical_weight * scale + title_weight = base_title_weight * scale # Normalize against the configured weight budget, not this query's adapted - # weight or whether its title arm happened to return rows, so scores stay + # weights or whether its title arm happened to return rows, so scores stay # comparable across the sub-searches Searcher merges (query + variants). weight_total = 1.0 + base_lexical_weight + ( - self._config.title_search_weight if self._config.title_search else 0.0 + base_title_weight if self._config.title_search else 0.0 ) fused = fuse_arms( vector_rows, self._fts_arm(table, query_text, top_k, chunk_type), title_rows, lexical_weight=lexical_weight, - title_weight=self._config.title_search_weight, + title_weight=title_weight, weight_total=weight_total, ) fused = _drop_unsupported_far_rows(fused, max_distance) diff --git a/src/lilbee/data/store/fusion.py b/src/lilbee/data/store/fusion.py index 66952a86d..743631f70 100644 --- a/src/lilbee/data/store/fusion.py +++ b/src/lilbee/data/store/fusion.py @@ -48,33 +48,38 @@ def vector_similarity(distance: float) -> float: return max(0.0, min(1.0, 1.0 - distance)) -def adaptive_lexical_weight( - vector_rows: list[SearchChunk], base_weight: float, margin_scale: float -) -> float: - """Shrink the lexical arm's weight toward zero when the vector arm is +def adaptive_weight_scale(vector_rows: list[SearchChunk], margin_scale: float) -> float: + """A [0, 1] factor to shrink the lexical arms by when the vector arm is confident about this query. A *peaked* vector ranking -- a top hit standing well clear of the field -- - means the dense embedder already located the answer and the lexical arm - mostly adds term-match noise. A *flat* ranking means dense is unsure and + means the dense embedder already located the answer and the lexical arms + mostly add term-match noise. A *flat* ranking means dense is unsure and BM25's exact-term matching is worth trusting. The confidence signal is the margin between the top similarity and the mean of the rest, divided by - *margin_scale*: at or above that margin the lexical arm is fully silenced, - at zero margin it keeps its full *base_weight*, and it scales linearly - between. Returns *base_weight* unchanged when there is nothing to measure - (fewer than two scored rows) or when *margin_scale* <= 0 (adaptation off). + *margin_scale*: at or above that margin the factor is 0 (arms silenced), at + zero margin it is 1 (arms kept), scaling linearly between. Returns 1.0 when + there is nothing to measure (fewer than two scored rows) or when + *margin_scale* <= 0 (adaptation off). """ if margin_scale <= 0: - return base_weight + return 1.0 sims = sorted( (vector_similarity(r.distance) for r in vector_rows if r.distance is not None), reverse=True, ) if len(sims) < _MIN_ROWS_FOR_MARGIN: - return base_weight + return 1.0 margin = max(0.0, sims[0] - fmean(sims[1:])) confidence = min(1.0, margin / margin_scale) - return base_weight * (1.0 - confidence) + return 1.0 - confidence + + +def adaptive_lexical_weight( + vector_rows: list[SearchChunk], base_weight: float, margin_scale: float +) -> float: + """The lexical arm's *base_weight* scaled by :func:`adaptive_weight_scale`.""" + return base_weight * adaptive_weight_scale(vector_rows, margin_scale) def normalized_bm25(scores: list[float]) -> list[float]: diff --git a/tests/test_query.py b/tests/test_query.py index cd300565e..023f8ed9c 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -27,6 +27,8 @@ unique_sources, ) from lilbee.retrieval.query.searcher import ( + _CONTEXT_TEMPLATE_TOKENS, + _PER_SOURCE_TOKENS, EMPTY_LIBRARY, GROUNDED_REFUSAL, SEARCH_NEEDS_EMBEDDER, @@ -992,6 +994,40 @@ def test_tight_budget_sheds_expansion_never_the_original(self, mock_svc): result = get_services().searcher.ask_raw("q") assert result.sources[0].chunk == "x" * 300 + def test_widened_prompt_still_fits_the_provider_ceiling(self, mock_svc): + """Widening runs after the budget fit, so it must budget against the same + provider ceiling: the assembled widened prompt must not exceed + prompt_token_budget, and expansion must actually happen (non-vacuous).""" + from lilbee.providers.base import prompt_token_budget + + ctx = 4096 + cfg.neighbor_expansion = 1 + cfg.num_ctx = ctx + mock_svc.provider.served_chat_ctx.return_value = None + mock_svc.store.get_chunks_by_indices.return_value = [ + _make_result(chunk="b" * 3000, chunk_index=1, page_start=2, page_end=2), + _make_result(chunk="a" * 3000, chunk_index=3, page_start=4, page_end=4), + ] + system, question = "sys " * 40, "q " * 20 + base = [_make_result(chunk="c" * 1500, chunk_index=2, page_start=3, page_end=3)] + + searcher = get_services().searcher + fitted = searcher._fit_context_budget(base, system, question, None, 1.0) + widened = searcher._widen_with_neighbors(fitted, system, question, None, 1.0) + + # Non-vacuous: at least one neighbor was actually merged in. + assert widened[0].chunk != "c" * 1500 + assembled = ( + searcher._budget_tokens(system) + + searcher._budget_tokens(question) + + sum(searcher._budget_tokens(r.chunk) + _PER_SOURCE_TOKENS for r in widened) + + _CONTEXT_TEMPLATE_TOKENS + ) + assert assembled <= prompt_token_budget(ctx), ( + f"widened prompt {assembled} tokens exceeds provider ceiling " + f"{prompt_token_budget(ctx)}" + ) + class TestAskRaw: def test_returns_structured_result(self, mock_svc): diff --git a/tests/test_store.py b/tests/test_store.py index 4ee7b1de7..1db569d73 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -180,8 +180,23 @@ def test_optimize_failure_keeps_hybrid_ready(self, store): store.ensure_fts_index() assert store._fts_ready is True - def test_first_call_creates_without_replace(self, store): - """Fresh table creates the chunk and title indexes, both with replace=False.""" + def test_first_call_creates_chunk_index_only_when_title_search_off(self, store): + """With title_search off (default), only the chunk index is built; the + title index is not created on a store that never queries it.""" + store.add_chunks(_make_records()) + table = store.open_table("chunks") + assert table is not None + + with mock.patch.object(type(table), "create_fts_index") as create_spy: + store.ensure_fts_index() + + assert [c.args[0] for c in create_spy.call_args_list] == ["chunk"] + assert create_spy.call_args_list[0].kwargs.get("with_position") is False + + def test_first_call_creates_both_indexes_when_title_search_on(self, store, test_config): + """Fresh table creates the chunk and title indexes, both positionless + with replace=False, when the title arm is enabled.""" + test_config.title_search = True store.add_chunks(_make_records()) table = store.open_table("chunks") assert table is not None @@ -690,10 +705,10 @@ def test_hybrid_search_with_fts_index(self, store, test_config): assert all(0.0 <= s <= 1.0 for s in scores) def test_adaptive_fusion_feeds_a_derived_weight_to_fusion(self, store, test_config): - """With adaptive_fusion on (the default), the per-query weight from - adaptive_lexical_weight -- fed the configured margin -- is what reaches - fuse_arms, not the fixed config value. Deleting the adaptive branch would - fail this, unlike a smoke test that only checks the score range.""" + """With adaptive_fusion on (the default), the per-query factor from + adaptive_weight_scale -- fed the configured margin -- scales the lexical + weight reaching fuse_arms, not the fixed config value. Deleting the + adaptive branch would fail this, unlike a smoke test on the score range.""" assert test_config.adaptive_fusion is True # shipped default test_config.adaptive_fusion_margin = 0.42 store.add_chunks(_make_records()) @@ -702,25 +717,25 @@ def test_adaptive_fusion_feeds_a_derived_weight_to_fusion(self, store, test_conf from lilbee.data.store import core as store_core with ( - mock.patch.object(store_core, "adaptive_lexical_weight", return_value=0.123) as adapt, + mock.patch.object(store_core, "adaptive_weight_scale", return_value=0.5) as scale, mock.patch.object(store_core, "fuse_arms", wraps=store_core.fuse_arms) as fuse, ): results = store.search(query_vec, top_k=3, query_text="chunk number") - adapt.assert_called_once() - # adaptive_lexical_weight(vector_rows, base_weight, margin): the base is - # lexical_fusion_weight and the margin is the configured value. - assert adapt.call_args.args[1] == pytest.approx(test_config.lexical_fusion_weight) - assert adapt.call_args.args[2] == pytest.approx(0.42) - assert fuse.call_args.kwargs["lexical_weight"] == pytest.approx(0.123) + scale.assert_called_once() + # adaptive_weight_scale(vector_rows, margin): the margin is the config value. + assert scale.call_args.args[1] == pytest.approx(0.42) + assert fuse.call_args.kwargs["lexical_weight"] == pytest.approx( + test_config.lexical_fusion_weight * 0.5 + ) # The normalization denominator is the configured budget (constant), not - # the adapted 0.123, so scores stay comparable across sub-searches. + # the adapted weight, so scores stay comparable across sub-searches. assert fuse.call_args.kwargs["weight_total"] == pytest.approx( 1.0 + test_config.lexical_fusion_weight ) assert len(results) > 0 def test_fixed_fusion_pins_the_config_weight(self, store, test_config): - """Opting out of adaptive fusion skips adaptive_lexical_weight and pins + """Opting out of adaptive fusion skips adaptive_weight_scale and pins the fixed lexical_fusion_weight into fuse_arms; a non-default weight must reach fusion verbatim.""" test_config.adaptive_fusion = False @@ -731,7 +746,7 @@ def test_fixed_fusion_pins_the_config_weight(self, store, test_config): from lilbee.data.store import core as store_core with ( - mock.patch.object(store_core, "adaptive_lexical_weight") as adapt, + mock.patch.object(store_core, "adaptive_weight_scale") as adapt, mock.patch.object(store_core, "fuse_arms", wraps=store_core.fuse_arms) as fuse, ): results = store.search(query_vec, top_k=3, query_text="chunk number") @@ -2466,6 +2481,28 @@ def test_title_search_weight_reaches_fusion(self, store, test_config): 1.0 + test_config.lexical_fusion_weight + 0.2 ) + def test_adaptive_fusion_scales_the_title_arm_too(self, store, test_config): + """Adaptive fusion downweights lexical; the title arm is also lexical, so + it must be scaled by the same confidence, not left at full weight (which + would re-admit the signal adaptive fusion just silenced).""" + test_config.title_search = True + test_config.title_search_weight = 0.5 + assert test_config.adaptive_fusion is True + store.add_chunks(_titled_records("a.pdf", 2, title="zebra manifesto")) + store.ensure_fts_index() + query_vec = [0.1] * test_config.embedding_dim + from lilbee.data.store import core as store_core + + with ( + mock.patch.object(store_core, "adaptive_weight_scale", return_value=0.25), + mock.patch.object(store_core, "fuse_arms", wraps=store_core.fuse_arms) as fuse, + ): + store.search(query_vec, top_k=4, max_distance=0, query_text="zebra") + assert fuse.call_args.kwargs["lexical_weight"] == pytest.approx( + test_config.lexical_fusion_weight * 0.25 + ) + assert fuse.call_args.kwargs["title_weight"] == pytest.approx(0.5 * 0.25) + def test_title_arm_off_by_default(self, store, test_config): """With title_search off, a title-only term earns no lexical support.""" assert test_config.title_search is False @@ -2478,6 +2515,7 @@ def test_title_arm_off_by_default(self, store, test_config): def test_title_arm_respects_chunk_type_filter(self, store, test_config): from lilbee.data.store.lance_helpers import _has_fts_index + test_config.title_search = True # the title index is built only when enabled store.add_chunks(_titled_records("a.pdf", 2, title="zebra manifesto")) store.ensure_fts_index() table = store.open_table("chunks") @@ -2485,6 +2523,19 @@ def test_title_arm_respects_chunk_type_filter(self, store, test_config): assert store._title_arm(table, "zebra", 5, ChunkType.RAW) assert store._title_arm(table, "zebra", 5, ChunkType.WIKI) == [] + def test_title_arm_failure_degrades_to_empty(self, store, test_config): + """A query-time title-arm failure returns no rows instead of raising, + so a broken title index can't take down the healthy chunk-BM25 arm and + collapse the whole hybrid search to vector-only.""" + test_config.title_search = True + store.add_chunks(_titled_records("a.pdf", 2, title="zebra manifesto")) + store.ensure_fts_index() + table = store.open_table("chunks") + from lilbee.data.store import core as store_core + + with mock.patch.object(store_core, "_lexical_rows", side_effect=RuntimeError("boom")): + assert store._title_arm(table, "zebra", 5, None) == [] + def test_old_index_without_title_column_still_searches(self, store, test_config): """Feature detection: a pre-title index searches fine and the title arm silently contributes nothing.""" @@ -2508,6 +2559,17 @@ def test_add_chunks_evolves_pre_title_table(self, store): rows = store.get_chunks_by_source("new.pdf") assert [r.title for r in rows] == ["fresh document"] + def test_pre_title_migration_warns_to_rebuild(self, store, caplog): + """Migrating an old store logs that pre-upgrade docs stay title-blind + until a rebuild, so the silence doesn't hide the gap.""" + import logging + + table = _create_pre_title_chunks_table(store) + table.add(_make_records()) + with caplog.at_level(logging.WARNING): + store.add_chunks(_titled_records("new.pdf", 1, title="fresh document")) + assert any("lilbee rebuild" in r.message for r in caplog.records) + def test_bm25_probe_stays_chunk_scoped(self, store): """The probe pins the chunk column: a title-only term is not a probe hit.""" store.add_chunks(_titled_records("a.pdf", 2, title="zebra manifesto")) From 1a37e7c2a8ccac1f2963a6de2503f45c96e335bd Mon Sep 17 00:00:00 2001 From: Tobias Perelstein <5562156+tobocop2@users.noreply.github.com> Date: Sun, 19 Jul 2026 20:40:13 -0400 Subject: [PATCH 16/28] Self-heal positional FTS indexes, expand ~ in data root, fix overflow-retry shedding Three more PR #557 audit findings: - A store whose FTS index was built with token positions overflows LanceDB's list encoding on every optimize() with no remediation short of a full re-ingest. That specific error is now caught and the index rebuilt positionless once, so index maintenance completes again; unrelated optimize failures still just keep the existing index. - LILBEE_DATA_ROOT is expanduser()'d, so a "~/lilbee" value from a systemd unit or .env points at the home directory instead of creating a literal "./~" tree (and diverging a path-keyed server lock). - The context-overflow retry refit from already-widened chunks, so it dropped whole lower-ranked original passages while keeping higher-ranked chunks' neighbor filler. RagContext now carries the pre-widen set and the retry refits from it, restoring "shed expansion before dropping an original." Server handlers and tests updated for the richer RagContext return type. Targeted suites green; full make check running. --- src/lilbee/core/config/model.py | 5 +- src/lilbee/data/store/core.py | 47 +++++++++++-- src/lilbee/retrieval/query/searcher.py | 30 ++++++--- src/lilbee/server/handlers/rag.py | 6 +- tests/server/test_chat_compaction.py | 4 +- tests/server/test_chat_stream_reasoning.py | 7 +- tests/server/test_rag_uses_chat_dispatch.py | 3 +- tests/test_config.py | 12 ++++ tests/test_query.py | 48 ++++++++++--- tests/test_server_handlers.py | 23 ++++--- tests/test_store.py | 75 +++++++++++++++++++++ 11 files changed, 218 insertions(+), 42 deletions(-) diff --git a/src/lilbee/core/config/model.py b/src/lilbee/core/config/model.py index a39f8bd21..287646158 100644 --- a/src/lilbee/core/config/model.py +++ b/src/lilbee/core/config/model.py @@ -1029,7 +1029,10 @@ def _resolve_defaults(cls, data: Any) -> Any: data["data_root"] = local if local is not None else default_data_dir() # data_root may arrive as a raw string (e.g. from LILBEE_DATA_ROOT); the # child-path derivations below use ``/``, so coerce to Path first. - root = Path(data["data_root"]) + # expanduser() so a "~/lilbee" value from a systemd unit or .env (which + # do not expand ~) points at the home dir instead of creating a literal + # ./~ tree that a server lock keyed on this path would then diverge on. + root = Path(data["data_root"]).expanduser() data["data_root"] = root if data.get("documents_dir") in (None, _UNSET_PATH): data["documents_dir"] = root / "documents" diff --git a/src/lilbee/data/store/core.py b/src/lilbee/data/store/core.py index 40570aad8..f3dddf2f4 100644 --- a/src/lilbee/data/store/core.py +++ b/src/lilbee/data/store/core.py @@ -108,6 +108,18 @@ def _drop_unsupported_far_rows( _MAX_FILTER_ITERATIONS = 20 # safety cap to prevent runaway loops +def _is_fts_position_overflow(exc: Exception) -> bool: + """True when *exc* is LanceDB's positional-FTS list-encoding overflow. + + A positional index (built by an intermediate dev commit) raises e.g. + "Max offset N exceeds length of values M" on optimize(); a positionless + rebuild is the remediation. Matched on message because LanceDB raises it + as a generic error type. + """ + msg = str(exc).lower() + return "offset" in msg and "exceeds" in msg + + def _lexical_rows( table: lancedb.table.Table, query_text: str, @@ -468,11 +480,18 @@ def ensure_fts_index(self) -> None: # table, the title index included. table.optimize() log.debug("FTS index optimized on '%s'", CHUNKS_TABLE) - except Exception: - log.warning( - "FTS optimize() failed; the existing index still serves hybrid search", - exc_info=True, - ) + except Exception as exc: + if _is_fts_position_overflow(exc): + # A pre-fix positional index overflows LanceDB's list + # encoding on every optimize(). Rebuild it positionless + # once so index maintenance can complete again. + self._rebuild_fts_positionless(table) + else: + log.warning( + "FTS optimize() failed; the existing index still " + "serves hybrid search", + exc_info=True, + ) else: # Positionless: with_position=True overflows LanceDB's list # encoding on optimize() for a large corpus. Positions only @@ -504,6 +523,24 @@ def _ensure_title_fts_unlocked(self, table: lancedb.table.Table) -> None: except Exception: log.debug("Title FTS index create failed", exc_info=True) + def _rebuild_fts_positionless(self, table: lancedb.table.Table) -> None: + """Replace positional FTS indexes with positionless ones. Caller holds the lock. + + The one-shot remediation for a store whose index was built + ``with_position=True`` and now overflows on every ``optimize()``. The + title index is rebuilt too when the title arm is enabled. + """ + try: + table.create_fts_index("chunk", replace=True, with_position=False) + if self._config.title_search and _TITLE_COLUMN in table.schema.names: + table.create_fts_index(_TITLE_COLUMN, replace=True, with_position=False) + log.warning("Rebuilt the FTS index positionless after a positional-index overflow") + except Exception: + log.warning( + "Positionless FTS rebuild failed; the existing index still serves", + exc_info=True, + ) + def ensure_scalar_indexes(self) -> None: """Build scalar indexes on the columns lilbee filters by. diff --git a/src/lilbee/retrieval/query/searcher.py b/src/lilbee/retrieval/query/searcher.py index d5b261942..234f0a509 100644 --- a/src/lilbee/retrieval/query/searcher.py +++ b/src/lilbee/retrieval/query/searcher.py @@ -227,10 +227,16 @@ class StructuredQuery(NamedTuple): class RagContext(NamedTuple): - """Grounded context for one turn: the chunks and the prompt built on them.""" + """Grounded context for one turn: the chunks and the prompt built on them. + + ``base_results`` is the pre-widen selected set. An overflow retry refits from + it, not from ``results`` (whose neighbor text is baked in and can no longer + be shed), so a tighter fit drops expansion before it drops an original chunk. + """ results: list[SearchChunk] messages: list[ChatMessage] + base_results: list[SearchChunk] | None = None class Searcher: @@ -931,6 +937,7 @@ def _finalize_context( retrieved set tighter without re-running retrieval or condensation. """ system = self._system_with_memory(self._config.rag_system_prompt, question) + base_results = list(results) results = self._fit_context_budget(results, system, question, history, scale) results = self._widen_with_neighbors(results, system, question, history, scale) context = build_context(results) @@ -939,7 +946,7 @@ def _finalize_context( if history: messages.extend(history) messages.append({"role": "user", "content": prompt}) - return RagContext(results, messages) + return RagContext(results, messages, base_results) @staticmethod def _budget_tokens(text: str) -> int: @@ -1310,7 +1317,7 @@ def ask_raw( rag = self.build_rag_context(question, top_k=top_k, history=history, chunk_type=chunk_type) if rag is None: return AskResult(answer=GROUNDED_REFUSAL, sources=[]) - results, messages = rag + results, messages = rag.results, rag.messages opts = options if options is not None else self._config.generation_options() try: result = self._provider.chat( @@ -1319,13 +1326,18 @@ def ask_raw( except ProviderError as exc: if exc.kind is not ProviderErrorKind.CONTEXT_OVERFLOW or not results: raise - # The budget estimator is a heuristic; when the engine still - # reports overflow, refit the same retrieved set tighter and - # retry once instead of hard-failing the question. + # The budget estimator is a heuristic; when the engine still reports + # overflow, refit tighter and retry once. Refit from the pre-widen + # set so the tighter budget sheds neighbor expansion before it drops + # an original chunk, not the reverse. log.warning("Context overflow despite budgeting; retrying with a tighter fit") - results, messages = self._finalize_context( - results, question, history, scale=_OVERFLOW_RETRY_SCALE + retry = self._finalize_context( + rag.base_results if rag.base_results is not None else results, + question, + history, + scale=_OVERFLOW_RETRY_SCALE, ) + results, messages = retry.results, retry.messages result = self._provider.chat( self._messages_for_provider(messages), options=opts or None ) @@ -1399,7 +1411,7 @@ def ask_stream( if rag is None: yield StreamToken(content=GROUNDED_REFUSAL, is_reasoning=False) return - results, messages = rag + results, messages = rag.results, rag.messages # No overflow retry here: a stream cannot be rebuilt once tokens have # been yielded, so the conservative budget in _fit_context_budget is # the streaming path's protection. diff --git a/src/lilbee/server/handlers/rag.py b/src/lilbee/server/handlers/rag.py index 470b2c315..25614bbad 100644 --- a/src/lilbee/server/handlers/rag.py +++ b/src/lilbee/server/handlers/rag.py @@ -509,7 +509,7 @@ async def chat( return AskResponse( answer=GROUNDED_REFUSAL, sources=[], cited_sources=[], compaction=compaction ) - sources, messages = rag + sources, messages = rag.results, rag.messages req = _build_canonical_request(messages, options) response = await asyncio.to_thread(dispatch_chat, req) text = _join_text_blocks(response.content) @@ -631,7 +631,7 @@ async def _stream_chat_response( yield frame if ctx is None: return - sources, messages = ctx + sources, messages = ctx.results, ctx.messages req = _build_canonical_request(messages, options) answer_parts: list[str] = [] @@ -852,7 +852,7 @@ def _resolve_stream_context( return _StreamResolution([], None, [frame]) if rag is None: return _StreamResolution([], None, [sse_error("No relevant documents found.")]) - results, messages = rag + results, messages = rag.results, rag.messages return _StreamResolution(results, messages, []) diff --git a/tests/server/test_chat_compaction.py b/tests/server/test_chat_compaction.py index 0f7c636f6..97e48723f 100644 --- a/tests/server/test_chat_compaction.py +++ b/tests/server/test_chat_compaction.py @@ -305,9 +305,11 @@ async def test_other_requests_are_served_while_retrieval_runs( released = False def _slow_context_signalling(question, top_k=0, history=None, chunk_type=None): + from lilbee.retrieval.query.searcher import RagContext + loop.call_soon_threadsafe(entered.set) time.sleep(0.3) # stands in for embed + search + rerank - return ([], [{"role": "user", "content": "q"}]) + return RagContext([], [{"role": "user", "content": "q"}]) mock_svc.searcher.build_rag_context.side_effect = _slow_context_signalling loop = asyncio.get_running_loop() diff --git a/tests/server/test_chat_stream_reasoning.py b/tests/server/test_chat_stream_reasoning.py index 42f85fa2f..bd05224f1 100644 --- a/tests/server/test_chat_stream_reasoning.py +++ b/tests/server/test_chat_stream_reasoning.py @@ -3,7 +3,6 @@ from __future__ import annotations from collections.abc import AsyncIterator -from typing import Any import pytest @@ -41,8 +40,10 @@ async def _gen(): return _gen() -def _rag_return() -> tuple[list[Any], list[dict[str, str]]]: - return ( +def _rag_return(): + from lilbee.retrieval.query.searcher import RagContext + + return RagContext( [], [ {"role": "system", "content": "ctx"}, diff --git a/tests/server/test_rag_uses_chat_dispatch.py b/tests/server/test_rag_uses_chat_dispatch.py index 6b90d6b63..c3740aa95 100644 --- a/tests/server/test_rag_uses_chat_dispatch.py +++ b/tests/server/test_rag_uses_chat_dispatch.py @@ -12,6 +12,7 @@ from lilbee.app.services import set_services from lilbee.core.config import cfg from lilbee.providers.base import ChatResult, FinishReason +from lilbee.retrieval.query.searcher import RagContext from lilbee.server import auth as _auth_mod from lilbee.server.chat_dispatch.canonical import ( CanonicalChatRequest, @@ -53,7 +54,7 @@ def services_with_chat_dispatch(): services = make_mock_services(provider=provider) services.registry.list_installed = MagicMock(return_value=[_installed_manifest(cfg.chat_model)]) services.searcher.build_rag_context = MagicMock( - return_value=( + return_value=RagContext( [], [ {"role": "system", "content": "ctx"}, diff --git a/tests/test_config.py b/tests/test_config.py index 7158e4241..1bfeaeb68 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -104,6 +104,18 @@ def test_lilbee_data_overrides_paths(self, tmp_path): assert c.data_dir == tmp_path / "data" assert c.lancedb_dir == tmp_path / "data" / "lancedb" + def test_data_root_expands_user_home(self): + """A ~ in LILBEE_DATA_ROOT expands: systemd/.env deliver a literal '~' + that would otherwise create a './~' tree and split a path-keyed lock.""" + env = _clean_env() + env.pop("LILBEE_DATA", None) + env["LILBEE_SKIP_TOML_CONFIG"] = "1" + env["LILBEE_DATA_ROOT"] = "~/lilbee_expanduser_probe" + with mock.patch.dict(os.environ, env, clear=True): + c = Config() + assert c.data_root == Path.home() / "lilbee_expanduser_probe" + assert c.lancedb_dir == Path.home() / "lilbee_expanduser_probe" / "data" / "lancedb" + def test_local_server_urls_from_env(self, tmp_path): env = _clean_env(tmp_path) env["LILBEE_OLLAMA_BASE_URL"] = "http://box:11434" diff --git a/tests/test_query.py b/tests/test_query.py index 023f8ed9c..92e7f0baa 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -994,6 +994,37 @@ def test_tight_budget_sheds_expansion_never_the_original(self, mock_svc): result = get_services().searcher.ask_raw("q") assert result.sources[0].chunk == "x" * 300 + def test_overflow_retry_sheds_expansion_not_an_original(self, mock_svc): + """On the context-overflow retry the tighter refit must start from the + pre-widen originals, so it sheds a chunk's neighbor filler before it + drops a whole lower-ranked original chunk.""" + from lilbee.providers.base import ProviderError, ProviderErrorKind + + cfg.neighbor_expansion = 1 + cfg.num_ctx = 4096 + mock_svc.provider.served_chat_ctx.return_value = None + a = _make_result(source="a.pdf", chunk="a" * 300, chunk_index=2, distance=0.1) + b = _make_result(source="b.pdf", chunk="b" * 300, chunk_index=0, distance=0.2) + mock_svc.store.search.return_value = [a, b] + + def neighbors(source, indices): + if source == "a.pdf": + return [ + _make_result(source="a.pdf", chunk="n" * 3000, chunk_index=1), + _make_result(source="a.pdf", chunk="m" * 3000, chunk_index=3), + ] + return [] + + mock_svc.store.get_chunks_by_indices.side_effect = neighbors + mock_svc.provider.chat.side_effect = [ + ProviderError("overflow", kind=ProviderErrorKind.CONTEXT_OVERFLOW), + _text_result("fits now"), + ] + result = get_services().searcher.ask_raw("q") + assert mock_svc.provider.chat.call_count == 2 + # The lower-ranked original b.pdf must survive; only a.pdf's expansion sheds. + assert {r.source for r in result.sources} == {"a.pdf", "b.pdf"} + def test_widened_prompt_still_fits_the_provider_ceiling(self, mock_svc): """Widening runs after the budget fit, so it must budget against the same provider ceiling: the assembled widened prompt must not exceed @@ -1024,8 +1055,7 @@ def test_widened_prompt_still_fits_the_provider_ceiling(self, mock_svc): + _CONTEXT_TEMPLATE_TOKENS ) assert assembled <= prompt_token_budget(ctx), ( - f"widened prompt {assembled} tokens exceeds provider ceiling " - f"{prompt_token_budget(ctx)}" + f"widened prompt {assembled} tokens exceeds provider ceiling {prompt_token_budget(ctx)}" ) @@ -2192,7 +2222,7 @@ def test_named_document_bypasses_similarity_search(self, mock_svc): ] rag = get_services().searcher.build_rag_context("summarize survey_report.pdf") assert rag is not None - results, _ = rag + results = rag.results assert [r.chunk for r in results] == ["first", "second"] assert all(r.score == 1.0 for r in results) mock_svc.store.search.assert_not_called() @@ -2239,7 +2269,7 @@ def test_routed_document_fits_the_served_context_window(self, mock_svc): cfg.num_ctx = None rag = get_services().searcher.build_rag_context("summarize survey_report.pdf") assert rag is not None - results, messages = rag + results, messages = rag.results, rag.messages prompt_tokens = sum(len(m["content"]) // 4 for m in messages) assert prompt_tokens <= 8192 # Document order preserved after trimming: the head survives. @@ -2259,7 +2289,7 @@ def test_dense_text_budgets_conservatively(self, mock_svc): cfg.num_ctx = None rag = get_services().searcher.build_rag_context("summarize survey_report.pdf") assert rag is not None - _, messages = rag + messages = rag.messages # At 2.56 chars/token the real prompt cost must still fit the window. real_tokens = sum(len(m["content"]) / 2.56 for m in messages) assert real_tokens <= 24576 * 1.05 @@ -2299,7 +2329,7 @@ def test_budget_falls_back_to_config_when_served_ctx_unknown(self, mock_svc): finally: cfg.num_ctx = None assert rag is not None - _, messages = rag + messages = rag.messages assert sum(len(m["content"]) // 4 for m in messages) <= 4096 def test_docket_reference_resolves_by_content_concentration(self, mock_svc): @@ -2570,7 +2600,7 @@ def test_follow_up_is_rewritten_for_retrieval(self, mock_svc): finally: cfg.query_expansion_count = 3 assert rag is not None - _, messages = rag + messages = rag.messages assert mock_svc.store.search.call_args[1]["query_text"] == rewritten assert "and when was it written?" in messages[-1]["content"] assert rewritten not in messages[-1]["content"] @@ -2974,7 +3004,7 @@ def test_build_rag_context_default_is_mixed_pool(self, mock_svc): mock_svc.store.search.return_value = [wiki_chunk, raw_chunk] result = get_services().searcher.build_rag_context("question") assert result is not None - chunks, _ = result + chunks = result.results assert len(chunks) == 2 def test_build_rag_context_forwards_chunk_type_to_store(self, mock_svc): @@ -3515,7 +3545,7 @@ def test_filters_high_distance_results(self, mock_svc): mock_svc.store.search.return_value = [close, far] result = get_services().searcher.build_rag_context("question") assert result is not None - results, _ = result + results = result.results sources = [r.source for r in results] assert "close.pdf" in sources assert "far.pdf" not in sources diff --git a/tests/test_server_handlers.py b/tests/test_server_handlers.py index 96365847e..f7a8cf064 100644 --- a/tests/test_server_handlers.py +++ b/tests/test_server_handlers.py @@ -13,6 +13,7 @@ from lilbee.data.ingest import SyncResult from lilbee.data.store import SearchChunk from lilbee.providers.base import ProviderError, ProviderErrorKind +from lilbee.retrieval.query.searcher import RagContext from lilbee.runtime.progress import SseErrorCode from lilbee.server import handlers from lilbee.server.handlers import ( @@ -42,9 +43,11 @@ def _rag_return(chunks: list[SearchChunk] | None = None): """Build a mock build_rag_context return value.""" + from lilbee.retrieval.query.searcher import RagContext + results = chunks or [_SAMPLE_CHUNK] messages = [{"role": "system", "content": "test"}, {"role": "user", "content": "q"}] - return results, messages + return RagContext(results, messages) @pytest.fixture(autouse=True) @@ -498,7 +501,7 @@ async def test_sources_event_falls_back_to_full_set_when_uncited(self, mock_svc) retrieved set, mirroring Searcher.ask_stream's ``used if used else results``.""" a = _SAMPLE_CHUNK.model_copy(update={"source": "a.md", "chunk_index": 0}) b = _SAMPLE_CHUNK.model_copy(update={"source": "b.md", "chunk_index": 1}) - mock_svc.searcher.build_rag_context.return_value = ([a, b], []) + mock_svc.searcher.build_rag_context.return_value = RagContext([a, b], []) mock_svc.provider.chat.return_value = iter(["an answer with no citation markers"]) events = [e async for e in handlers.ask_stream("question")] sources_event = next(e for e in events if e and e.startswith("event: sources")) @@ -523,7 +526,7 @@ async def test_sources_event_carries_cited_subset(self, mock_svc): matching the non-stream cited_sources contract, not the full retrieved list.""" cited = _SAMPLE_CHUNK.model_copy(update={"source": "cited.md", "chunk_index": 0}) other = _SAMPLE_CHUNK.model_copy(update={"source": "other.md", "chunk_index": 1}) - mock_svc.searcher.build_rag_context.return_value = ([cited, other], []) + mock_svc.searcher.build_rag_context.return_value = RagContext([cited, other], []) mock_svc.provider.chat.return_value = iter(["see [1] for details"]) events = [e async for e in handlers.ask_stream("question")] sources_event = next(e for e in events if e and e.startswith("event: sources")) @@ -534,7 +537,7 @@ async def test_model_sources_block_suppressed_in_grounded_stream(self, mock_svc) """A model that appends its own Sources block must not double up with the authoritative SOURCES event: no token frame carries the model's list.""" cited = _SAMPLE_CHUNK.model_copy(update={"source": "real.md", "chunk_index": 0}) - mock_svc.searcher.build_rag_context.return_value = ([cited], []) + mock_svc.searcher.build_rag_context.return_value = RagContext([cited], []) mock_svc.provider.chat.return_value = iter( ["Grounded answer [1].", "\n\nSources:\n- invented.md"] ) @@ -555,7 +558,7 @@ async def test_grounded_stream_releases_held_back_final_line(self, mock_svc): """The citation filter holds the last line until the stream ends; the flushed tail must still be emitted as a token before SOURCES.""" cited = _SAMPLE_CHUNK.model_copy(update={"source": "real.md", "chunk_index": 0}) - mock_svc.searcher.build_rag_context.return_value = ([cited], []) + mock_svc.searcher.build_rag_context.return_value = RagContext([cited], []) mock_svc.provider.chat.return_value = iter(["First line [1].\n", "Held final line."]) events = [e async for e in handlers.ask_stream("question")] token_text = "".join( @@ -570,7 +573,7 @@ async def test_grounded_stream_forwards_reasoning_tokens(self, mock_svc, monkeyp channel and stays out of the answer (and the citation filter).""" monkeypatch.setattr(cfg, "show_reasoning", True) cited = _SAMPLE_CHUNK.model_copy(update={"source": "real.md", "chunk_index": 0}) - mock_svc.searcher.build_rag_context.return_value = ([cited], []) + mock_svc.searcher.build_rag_context.return_value = RagContext([cited], []) mock_svc.provider.chat.return_value = iter(["pondering", "answer [1]"]) events = [e async for e in handlers.ask_stream("question")] reasoning_text = "".join( @@ -1028,7 +1031,7 @@ async def test_sources_event_carries_cited_subset(self, mock_svc, monkeypatch): matching the non-stream cited_sources contract, not the full retrieved list.""" cited = _SAMPLE_CHUNK.model_copy(update={"source": "cited.md", "chunk_index": 0}) other = _SAMPLE_CHUNK.model_copy(update={"source": "other.md", "chunk_index": 1}) - mock_svc.searcher.build_rag_context.return_value = ([cited, other], []) + mock_svc.searcher.build_rag_context.return_value = RagContext([cited, other], []) monkeypatch.setattr( _rag_h, "dispatch_chat_stream", lambda req: _canonical_text_stream(["see [1] here"]) ) @@ -1042,7 +1045,7 @@ async def test_sources_event_falls_back_to_full_set_when_uncited(self, mock_svc, retrieved set, mirroring Searcher.ask_stream's ``used if used else results``.""" a = _SAMPLE_CHUNK.model_copy(update={"source": "a.md", "chunk_index": 0}) b = _SAMPLE_CHUNK.model_copy(update={"source": "b.md", "chunk_index": 1}) - mock_svc.searcher.build_rag_context.return_value = ([a, b], []) + mock_svc.searcher.build_rag_context.return_value = RagContext([a, b], []) monkeypatch.setattr( _rag_h, "dispatch_chat_stream", lambda req: _canonical_text_stream(["no markers here"]) ) @@ -1055,7 +1058,7 @@ async def test_model_sources_block_suppressed_in_grounded_stream(self, mock_svc, """A model Sources block in chat streaming is dropped so it doesn't double up with the authoritative SOURCES event.""" cited = _SAMPLE_CHUNK.model_copy(update={"source": "real.md", "chunk_index": 0}) - mock_svc.searcher.build_rag_context.return_value = ([cited], []) + mock_svc.searcher.build_rag_context.return_value = RagContext([cited], []) monkeypatch.setattr( _rag_h, "dispatch_chat_stream", @@ -1077,7 +1080,7 @@ async def test_model_sources_block_suppressed_in_grounded_stream(self, mock_svc, async def test_chat_stream_releases_held_back_final_line(self, mock_svc, monkeypatch): """The chat SSE path also flushes the filter's held-back tail as a token.""" cited = _SAMPLE_CHUNK.model_copy(update={"source": "real.md", "chunk_index": 0}) - mock_svc.searcher.build_rag_context.return_value = ([cited], []) + mock_svc.searcher.build_rag_context.return_value = RagContext([cited], []) monkeypatch.setattr( _rag_h, "dispatch_chat_stream", diff --git a/tests/test_store.py b/tests/test_store.py index 1db569d73..d9d2d5050 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -180,6 +180,81 @@ def test_optimize_failure_keeps_hybrid_ready(self, store): store.ensure_fts_index() assert store._fts_ready is True + def test_positional_index_overflow_rebuilds_positionless(self, store): + """A store whose FTS index was built with positions overflows on every + optimize(); catching that specific error rebuilds the index positionless + (replace=True) so index maintenance can complete instead of failing + forever with no remediation short of a full re-ingest.""" + store.add_chunks(_make_records()) + store.ensure_fts_index() + table = store.open_table("chunks") + assert table is not None + overflow = RuntimeError("Max offset 1897296 exceeds length of values 1067891") + with ( + mock.patch.object(type(table), "optimize", side_effect=overflow), + mock.patch.object(type(table), "create_fts_index") as rebuild, + ): + store.ensure_fts_index() + assert any( + c.args[:1] == ("chunk",) + and c.kwargs.get("replace") is True + and c.kwargs.get("with_position") is False + for c in rebuild.call_args_list + ) + assert store._fts_ready is True + + def test_generic_optimize_failure_does_not_rebuild(self, store): + """An unrelated optimize() failure keeps the existing index and does NOT + pay for a full positionless rebuild it cannot fix.""" + store.add_chunks(_make_records()) + store.ensure_fts_index() + table = store.open_table("chunks") + assert table is not None + with ( + mock.patch.object(type(table), "optimize", side_effect=RuntimeError("disk full")), + mock.patch.object(type(table), "create_fts_index") as rebuild, + ): + store.ensure_fts_index() + rebuild.assert_not_called() + + def test_overflow_rebuild_also_rebuilds_the_title_index_when_enabled(self, store, test_config): + """With the title arm on, a positional-overflow rebuild replaces the + title index too, not just the chunk index.""" + test_config.title_search = True + store.add_chunks(_titled_records("a.pdf", 2, title="zebra manifesto")) + store.ensure_fts_index() + table = store.open_table("chunks") + assert table is not None + overflow = RuntimeError("Max offset 9 exceeds length of values 3") + with ( + mock.patch.object(type(table), "optimize", side_effect=overflow), + mock.patch.object(type(table), "create_fts_index") as rebuild, + ): + store.ensure_fts_index() + rebuilt = { + c.args[0] + for c in rebuild.call_args_list + if c.kwargs.get("replace") is True and c.kwargs.get("with_position") is False + } + assert rebuilt == {"chunk", "title"} + + def test_overflow_rebuild_failure_is_swallowed(self, store): + """If the positionless rebuild itself fails, the store keeps the existing + index and does not propagate: a failed self-heal must not crash sync.""" + store.add_chunks(_make_records()) + store.ensure_fts_index() + table = store.open_table("chunks") + assert table is not None + overflow = RuntimeError("Max offset 9 exceeds length of values 3") + with ( + mock.patch.object(type(table), "optimize", side_effect=overflow), + mock.patch.object( + type(table), "create_fts_index", side_effect=RuntimeError("rebuild boom") + ), + ): + store.ensure_fts_index() # must not raise + assert store._fts_ready is True + def test_first_call_creates_chunk_index_only_when_title_search_off(self, store): """With title_search off (default), only the chunk index is built; the title index is not created on a store that never queries it.""" From a75d9e8a2f577b11188b4c98c967ef2c8bbdf84c Mon Sep 17 00:00:00 2001 From: Tobias Perelstein <5562156+tobocop2@users.noreply.github.com> Date: Sun, 19 Jul 2026 22:02:57 -0400 Subject: [PATCH 17/28] Harden ingest metadata and config edges; close audit test gaps Verified residue leads from the PR #557 branch audit that reproduced on tip: - source_meta_from_extraction no longer trusts untyped extractor metadata: a string "authors" is treated as one author instead of being iterated into its characters, non-string entries are coerced, and a non-string title falls back to the filename stem instead of raising mid-ingest. - An empty LILBEE_DATA_ROOT ("") now falls through to default resolution like an unset one, instead of resolving the data root to the process cwd. - Structural detector: split the TOC line-count floor from the dot-leader hit floor (one constant did both jobs), and reword a fusion-weight comment off "historical behaviour" phrasing. - Tests: assert the optimize-failure warning is actually logged; assert the structural filter runs before the top_k*2 slice so real passages backfill; cover the four adaptive-fusion / structural-filter settings-map entries. make check pending; targeted suites green. --- src/lilbee/core/config/model.py | 6 +++++- src/lilbee/data/ingest/title.py | 14 +++++++++++--- src/lilbee/retrieval/query/structural.py | 4 +++- tests/test_config.py | 13 +++++++++++++ tests/test_ingest_title.py | 15 +++++++++++++++ tests/test_query.py | 21 +++++++++++++++++++++ tests/test_settings.py | 19 +++++++++++++++++++ tests/test_store.py | 9 +++++++-- tests/test_store_fusion.py | 2 +- 9 files changed, 95 insertions(+), 8 deletions(-) diff --git a/src/lilbee/core/config/model.py b/src/lilbee/core/config/model.py index 287646158..f0852ca24 100644 --- a/src/lilbee/core/config/model.py +++ b/src/lilbee/core/config/model.py @@ -205,7 +205,7 @@ class Config(BaseSettings): title_search_weight: float = ConfigField(default=0.5, ge=0.0, le=1.0, writable=True) # Lexical (BM25) arm weight relative to the vector arm in rank fusion. - # 1.0 keeps the two arms equal (the historical behaviour); lowering it lets + # 1.0 gives the two arms equal voice; lowering it lets # a strong dense embedder dominate on corpora where the lexical arm adds # noise rather than signal. The right value is corpus-dependent and set by # the retrieval benchmark, not guessed here. @@ -1020,6 +1020,10 @@ def _resolve_defaults(cls, data: Any) -> Any: if not isinstance(data, dict): return data + # An empty LILBEE_DATA_ROOT (delivered as "") must fall through to default + # resolution like an unset one, not become Path(".") = the process cwd. + if isinstance(data.get("data_root"), str) and not data["data_root"].strip(): + data["data_root"] = None if data.get("data_root") in (None, _UNSET_PATH): data_env = os.environ.get("LILBEE_DATA", "").strip() if data_env: diff --git a/src/lilbee/data/ingest/title.py b/src/lilbee/data/ingest/title.py index c34003002..af34b16ed 100644 --- a/src/lilbee/data/ingest/title.py +++ b/src/lilbee/data/ingest/title.py @@ -19,7 +19,7 @@ def derive_title(source_name: str, metadata_title: str | None = None) -> str: The stem cleanup flattens underscore/hyphen separators to spaces so BM25 tokenizes ``survey_214.pdf`` into the same terms a query would use. """ - if metadata_title and metadata_title.strip(): + if isinstance(metadata_title, str) and metadata_title.strip(): return metadata_title.strip() return _STEM_SEPARATOR_RE.sub(" ", PurePath(source_name).stem).strip() @@ -28,9 +28,17 @@ def source_meta_from_extraction(metadata: Mapping[str, Any], source_name: str) - """Fold kreuzberg extraction metadata into a :class:`SourceMeta`. The title falls back to the filename stem; authors and creation date stay - empty (persisted NULL) when the extractor reports none. + empty (persisted NULL) when the extractor reports none. Extractor metadata is + untyped, so a string ``authors`` is treated as one author (not split into its + characters) and non-string entries are coerced. """ - authors = metadata.get("authors") or [] + raw_authors = metadata.get("authors") + if isinstance(raw_authors, str): + authors: list[str] = [raw_authors] + elif isinstance(raw_authors, (list, tuple)): + authors = [str(a) for a in raw_authors if a] + else: + authors = [] return SourceMeta( title=derive_title(source_name, metadata.get("title")), authors=", ".join(a for a in authors if a), diff --git a/src/lilbee/retrieval/query/structural.py b/src/lilbee/retrieval/query/structural.py index 493f56bd6..ed59da2db 100644 --- a/src/lilbee/retrieval/query/structural.py +++ b/src/lilbee/retrieval/query/structural.py @@ -21,6 +21,8 @@ # A chunk needs at least this many non-empty lines before the TOC ratio means anything. _MIN_TOC_LINES = 3 +# And at least this many dot-leader lines, so a page with one stray "... 42" is not a TOC. +_MIN_TOC_HITS = 3 _TOC_RATIO = 0.30 # Cover/title-page gates: a title page is very short with essentially no prose. @@ -37,7 +39,7 @@ def _is_toc(nonempty: list[str]) -> bool: if len(nonempty) < _MIN_TOC_LINES: return False hits = sum(1 for line in nonempty if _TOC_LINE.search(line)) - return hits >= _MIN_TOC_LINES and hits / len(nonempty) >= _TOC_RATIO + return hits >= _MIN_TOC_HITS and hits / len(nonempty) >= _TOC_RATIO def _is_cover_page(text: str) -> bool: diff --git a/tests/test_config.py b/tests/test_config.py index 1bfeaeb68..60fc5914b 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -116,6 +116,19 @@ def test_data_root_expands_user_home(self): assert c.data_root == Path.home() / "lilbee_expanduser_probe" assert c.lancedb_dir == Path.home() / "lilbee_expanduser_probe" / "data" / "lancedb" + def test_empty_data_root_falls_back_to_default_not_cwd(self): + """An empty LILBEE_DATA_ROOT must resolve to the platform default, not + the process cwd (which would make the data dir move with the launcher).""" + env = _clean_env() + env.pop("LILBEE_DATA", None) + env["LILBEE_SKIP_TOML_CONFIG"] = "1" + env["LILBEE_DATA_ROOT"] = "" + with mock.patch.dict(os.environ, env, clear=True): + c = Config() + assert c.data_root != Path() + assert c.data_root != Path.cwd() + assert str(c.data_root).endswith("lilbee") + def test_local_server_urls_from_env(self, tmp_path): env = _clean_env(tmp_path) env["LILBEE_OLLAMA_BASE_URL"] = "http://box:11434" diff --git a/tests/test_ingest_title.py b/tests/test_ingest_title.py index a3ebccbf8..b29182eed 100644 --- a/tests/test_ingest_title.py +++ b/tests/test_ingest_title.py @@ -55,3 +55,18 @@ def test_none_values_tolerated(self): assert meta.title == "notes" assert meta.authors == "" assert meta.created_at == "" + + def test_string_authors_is_one_author_not_split_into_characters(self): + # A raw PDF /Author field often arrives as a plain string; it must not + # be iterated into "J, o, h, n". + meta = source_meta_from_extraction({"authors": "John Doe"}, "x.pdf") + assert meta.authors == "John Doe" + + def test_non_string_author_entries_are_coerced_not_raised(self): + meta = source_meta_from_extraction({"authors": ["Ada", 42]}, "x.pdf") + assert meta.authors == "Ada, 42" + + def test_non_string_title_falls_back_to_stem(self): + # A bytes/number title in malformed metadata must not raise; fall back. + meta = source_meta_from_extraction({"title": 123}, "annual_report.pdf") + assert meta.title == "annual report" diff --git a/tests/test_query.py b/tests/test_query.py index 92e7f0baa..6ba790ed8 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -515,6 +515,27 @@ def test_structural_chunk_the_query_missed_is_dropped(self, mock_svc): results = get_services().searcher.search("q") assert [r.source for r in results] == ["body.pdf"] + def test_structural_filter_runs_before_the_buffer_slice_so_real_passages_backfill( + self, mock_svc + ): + """The filter drops structural chunks from the candidate list before the + top_k*2 slice, so a real passage past the window backfills into it. If the + filter ran after the slice, the deeper real passage would be lost.""" + old_top_k = cfg.top_k + cfg.top_k = 1 # window is top_k*2 = 2 + cfg.filter_structural_chunks = True + try: + real0 = _make_result(source="a.pdf", distance=0.1, chunk="Answer prose one.") + toc1 = _make_result(source="toc1.pdf", distance=0.2, chunk=self._TOC_CHUNK) + toc2 = _make_result(source="toc2.pdf", distance=0.3, chunk=self._TOC_CHUNK) + real3 = _make_result(source="b.pdf", distance=0.4, chunk="Answer prose two.") + mock_svc.store.search.return_value = [real0, toc1, toc2, real3] + sources = [r.source for r in get_services().searcher.search("q")] + finally: + cfg.top_k = old_top_k + assert "b.pdf" in sources # backfilled from past the pre-filter window + assert "toc1.pdf" not in sources and "toc2.pdf" not in sources + def test_structural_filter_keeps_query_matched_page(self, mock_svc): """A page the query lexically hit is never dropped, whatever its shape: it is content the answer may need.""" diff --git a/tests/test_settings.py b/tests/test_settings.py index 021dc02db..615e05857 100644 --- a/tests/test_settings.py +++ b/tests/test_settings.py @@ -222,6 +222,25 @@ def test_neighbor_expansion_in_settings_map(self): assert defn.group == "Retrieval" assert get_default("neighbor_expansion") == 0 + def test_fusion_knobs_in_settings_map(self): + """The four adaptive-fusion / structural-filter knobs (which gate the + on-by-default fusion behavior) are on the settings surface with their + shipped defaults, so a dropped or typo'd entry fails CI.""" + from lilbee.app.settings_map import SETTINGS_MAP, get_default + + assert get_default("lexical_fusion_weight") == 1.0 + assert get_default("adaptive_fusion") is True + assert get_default("adaptive_fusion_margin") == 0.15 + assert get_default("filter_structural_chunks") is False + for key in ( + "lexical_fusion_weight", + "adaptive_fusion", + "adaptive_fusion_margin", + "filter_structural_chunks", + ): + assert SETTINGS_MAP[key].writable is True, key + assert SETTINGS_MAP[key].group == "Retrieval", key + 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 d9d2d5050..baa1d088d 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -203,9 +203,12 @@ def test_positional_index_overflow_rebuilds_positionless(self, store): ) assert store._fts_ready is True - def test_generic_optimize_failure_does_not_rebuild(self, store): + def test_generic_optimize_failure_does_not_rebuild(self, store, caplog): """An unrelated optimize() failure keeps the existing index and does NOT - pay for a full positionless rebuild it cannot fix.""" + pay for a full positionless rebuild it cannot fix. It must also log a + warning so an operator debugging a large corpus is not left in silence.""" + import logging + store.add_chunks(_make_records()) store.ensure_fts_index() table = store.open_table("chunks") @@ -213,9 +216,11 @@ def test_generic_optimize_failure_does_not_rebuild(self, store): with ( mock.patch.object(type(table), "optimize", side_effect=RuntimeError("disk full")), mock.patch.object(type(table), "create_fts_index") as rebuild, + caplog.at_level(logging.WARNING), ): store.ensure_fts_index() rebuild.assert_not_called() + assert any("optimize()" in r.message for r in caplog.records) def test_overflow_rebuild_also_rebuilds_the_title_index_when_enabled(self, store, test_config): """With the title arm on, a positional-overflow rebuild replaces the diff --git a/tests/test_store_fusion.py b/tests/test_store_fusion.py index 7721378c7..28035ced8 100644 --- a/tests/test_store_fusion.py +++ b/tests/test_store_fusion.py @@ -204,7 +204,7 @@ def _lex_row(self, weight: float): lex = [_chunk("cat.pdf", 0, bm25=35.0)] return next(r for r in fuse_arms(vec, lex, lexical_weight=weight) if r.source == "cat.pdf") - def test_default_weight_is_the_historical_equal_voice(self): + def test_default_weight_gives_the_arms_equal_voice(self): vec = [_chunk("noise.md", 0, distance=0.5)] lex = [_chunk("cat.pdf", 0, bm25=35.0)] default = {r.source: r.score for r in fuse_arms(vec, lex)} From d9440cdd29c886827ad7bde29021937f7cb747f5 Mon Sep 17 00:00:00 2001 From: Tobias Perelstein <5562156+tobocop2@users.noreply.github.com> Date: Sun, 19 Jul 2026 22:39:38 -0400 Subject: [PATCH 18/28] Fix the audit's merge-blocking P3 cluster (title/scalar/config) Addresses bb-z538, bb-t1jz, bb-dpo6, bb-gxws, bb-9f3o, bb-m9ok on feat/retrieval-consolidated: - Title FTS index-creation failure now warns instead of logging at debug: the new gating makes that branch fire only when title_search is enabled, so a silent failure meant an opted-in arm that quietly did nothing. - Scalar indexes now build lazily on the serve/search path (guarded by a _scalar_ready flag reset on writes), so a store served without a fresh ingest stops paying full-table prefilter scans. - ensure_scalar_indexes gives each column its own try (a BTree failure no longer skips the Bitmap) and warns on a populated table instead of a debug-only swallow with a guessed cause. - Also index chunk_concepts.chunk_source: the concept-boost path filters it once per search result, and a main-side boost_results fix depends on the index existing. - Chunk-level title persists NULL (not "") for an empty title, on both the ingest and import paths, matching the migration and the _sources table. - Added bounds-rejection tests for the four new numeric config fields, and an le=100 upper bound on neighbor_expansion so a token-count misread can't build a megabyte-long IN-predicate per query. --- src/lilbee/core/config/model.py | 6 ++- src/lilbee/data/export.py | 2 +- src/lilbee/data/ingest/pipeline.py | 4 +- src/lilbee/data/ingest/types.py | 5 +- src/lilbee/data/store/core.py | 71 +++++++++++++++++++++++------ tests/test_config.py | 24 ++++++++++ tests/test_store.py | 73 ++++++++++++++++++++++++++++-- 7 files changed, 163 insertions(+), 22 deletions(-) diff --git a/src/lilbee/core/config/model.py b/src/lilbee/core/config/model.py index f0852ca24..104cbb4f1 100644 --- a/src/lilbee/core/config/model.py +++ b/src/lilbee/core/config/model.py @@ -283,7 +283,11 @@ class Config(BaseSettings): # 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) + # Capped: it is a small chunk radius (useful values are single digits), and + # the merged text is token-budget-bounded anyway, so a large value only + # inflates per-query fetch cost -- and a misread as a token count (e.g. + # 50000) would build a megabyte-long IN-predicate per source. + neighbor_expansion: int = ConfigField(default=0, ge=0, le=100, 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/export.py b/src/lilbee/data/export.py index 5a2703374..ab541960a 100644 --- a/src/lilbee/data/export.py +++ b/src/lilbee/data/export.py @@ -230,7 +230,7 @@ async def import_dataset( # imported chunks visible to the title search arm. title = derive_title(name) for chunk in chunks: - chunk["title"] = title + chunk["title"] = title or None # One locked transaction (cleanup + chunks + page texts + source row) so a # failure can't leave the source with its old rows deleted and no new ones; # the embedding-dim check inside runs before the cleanup delete. diff --git a/src/lilbee/data/ingest/pipeline.py b/src/lilbee/data/ingest/pipeline.py index df702b4f5..1daf2870b 100644 --- a/src/lilbee/data/ingest/pipeline.py +++ b/src/lilbee/data/ingest/pipeline.py @@ -233,7 +233,9 @@ async def _produce_records( meta = extracted_meta[0] for record in records: - record["title"] = meta.title + # NULL (not "") for an absent title, so chunk rows match the migration + # and the _sources table, which both persist absence as NULL. + record["title"] = meta.title or None if meta_out is not None: meta_out.append(meta) return records diff --git a/src/lilbee/data/ingest/types.py b/src/lilbee/data/ingest/types.py index f66999efc..83f02278c 100644 --- a/src/lilbee/data/ingest/types.py +++ b/src/lilbee/data/ingest/types.py @@ -73,8 +73,9 @@ class ChunkRecord(TypedDict): chunk: str chunk_index: int vector: list[float] - # Stamped once per document by the pipeline (see _produce_records). - title: NotRequired[str] + # Stamped once per document by the pipeline (see _produce_records); None + # when the title is empty, so chunk rows persist NULL like the _sources table. + title: NotRequired[str | None] class SyncResult(BaseModel): diff --git a/src/lilbee/data/store/core.py b/src/lilbee/data/store/core.py index f3dddf2f4..8488c72d9 100644 --- a/src/lilbee/data/store/core.py +++ b/src/lilbee/data/store/core.py @@ -219,6 +219,10 @@ class Store: def __init__(self, config: Config) -> None: self._config = config self._fts_ready: bool = False + # Scalar indexes (source/chunk_type) are built at ingest, but a store + # served without a fresh ingest never ran that path, so search builds + # them lazily once. This guards the one-shot per process. + self._scalar_ready: bool = False self._db: lancedb.DBConnection | None = None # Cache of {filename: ingested_at} rebuilt only when sources # mutate; callers (temporal filter) hit it per-query. @@ -521,7 +525,13 @@ def _ensure_title_fts_unlocked(self, table: lancedb.table.Table) -> None: table.create_fts_index(_TITLE_COLUMN, replace=False, with_position=False) log.debug("Title FTS index created on '%s'", CHUNKS_TABLE) except Exception: - log.debug("Title FTS index create failed", exc_info=True) + # Only reached with title_search enabled, so a silent failure means + # the user's opted-in title arm quietly does nothing. Warn, don't hide. + log.warning( + "Title FTS index creation failed; the title-search arm will " + "contribute nothing until it can be built", + exc_info=True, + ) def _rebuild_fts_positionless(self, table: lancedb.table.Table) -> None: """Replace positional FTS indexes with positionless ones. Caller holds the lock. @@ -552,19 +562,43 @@ def ensure_scalar_indexes(self) -> None: skipped; ``optimize()`` folds later rows into every index, this one too. """ with self._write_lock(): - table = self.open_table(CHUNKS_TABLE) - if table is None: - return - names = table.schema.names + self._ensure_scalar_index_on( + CHUNKS_TABLE, (("source", "BTREE"), ("chunk_type", "BITMAP")) + ) + # The concept-boost path filters chunk_concepts by chunk_source once + # per search result (see ConceptGraph._chunk_concepts_from), so index + # it too or every boosted query full-scans the table. + self._ensure_scalar_index_on(CHUNK_CONCEPTS_TABLE, (("chunk_source", "BTREE"),)) + self._scalar_ready = True + + def _ensure_scalar_index_on( + self, table_name: str, columns: tuple[tuple[str, str], ...] + ) -> None: + """Build the given (column, index_type) scalar indexes on *table_name*. + + Caller holds ``write_lock()``. Each column gets its own try so one + failure does not skip the rest; a failure on a populated table warns + (the prefilter speedup is silently lost) while an empty table's is debug. + """ + table = self.open_table(table_name) + if table is None: + return + names = table.schema.names + fail_level = logging.WARNING if table.count_rows() > 0 else logging.DEBUG + for column, index_type in columns: + if column not in names or _has_scalar_index(table, column): + continue try: - if "source" in names and not _has_scalar_index(table, "source"): - table.create_scalar_index("source", index_type="BTREE", replace=False) - log.debug("Scalar (BTree) index created on 'source'") - if "chunk_type" in names and not _has_scalar_index(table, "chunk_type"): - table.create_scalar_index("chunk_type", index_type="BITMAP", replace=False) - log.debug("Scalar (Bitmap) index created on 'chunk_type'") + table.create_scalar_index(column, index_type=index_type, replace=False) + log.debug("Scalar (%s) index created on '%s.%s'", index_type, table_name, column) except Exception: - log.debug("Scalar index ensure failed (empty table?)", exc_info=True) + log.log( + fail_level, + "Scalar index create failed on '%s.%s'", + table_name, + column, + exc_info=True, + ) def ensure_vector_index(self, *, force: bool = False) -> bool: """Build or refresh the ANN vector index when the corpus is large enough. @@ -617,6 +651,7 @@ def add_chunks(self, records: list[dict]) -> int: embedding_dim = self._config.embedding_dim self._ensure_embedding_compat() self._fts_ready = False + self._scalar_ready = False if not records: return 0 _check_vector_dims(records, embedding_dim) @@ -680,6 +715,11 @@ def search( self.canonicalize_meta_if_legacy() self._ensure_embedding_compat() + if not self._scalar_ready: + # A serve-only store never ran ingest, where scalar indexes are + # built; without them the source/chunk_type prefilters full-scan. + self.ensure_scalar_indexes() + if query_text and not self._fts_ready: self.ensure_fts_index() @@ -813,8 +853,8 @@ def _hybrid_search( # Normalize against the configured weight budget, not this query's adapted # weights or whether its title arm happened to return rows, so scores stay # comparable across the sub-searches Searcher merges (query + variants). - weight_total = 1.0 + base_lexical_weight + ( - base_title_weight if self._config.title_search else 0.0 + weight_total = ( + 1.0 + base_lexical_weight + (base_title_weight if self._config.title_search else 0.0) ) fused = fuse_arms( vector_rows, @@ -1293,6 +1333,7 @@ def write_chunks_batch(self, items: list[ChunkWrite]) -> int: embedding_dim = self._config.embedding_dim self._ensure_embedding_compat() self._fts_ready = False + self._scalar_ready = False all_records = [rec for it in items for rec in it.records] _check_vector_dims(all_records, embedding_dim) db = self.get_db() @@ -1650,6 +1691,7 @@ def close(self) -> None: """Release the database connection and reset state.""" self._db = None self._fts_ready = False + self._scalar_ready = False def drop_all(self) -> None: """Drop every table except ``_memories`` -- used by rebuild. @@ -1660,6 +1702,7 @@ def drop_all(self) -> None: """ with self._write_lock(): self._fts_ready = False + self._scalar_ready = False db = self.get_db() for name in _table_names(db): if name == MEMORIES_TABLE: diff --git a/tests/test_config.py b/tests/test_config.py index 60fc5914b..59257d58a 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -128,6 +128,12 @@ def test_empty_data_root_falls_back_to_default_not_cwd(self): assert c.data_root != Path() assert c.data_root != Path.cwd() assert str(c.data_root).endswith("lilbee") + # A raw blank string (direct construction / a string env source) hits + # the same fall-through instead of resolving to Path(".") = cwd. + c2 = Config(data_root=" ") + assert c2.data_root != Path() + assert c2.data_root != Path.cwd() + assert str(c2.data_root).endswith("lilbee") def test_local_server_urls_from_env(self, tmp_path): env = _clean_env(tmp_path) @@ -210,6 +216,24 @@ def test_normalize_model_tag_blank_chat_rejected(self): with pytest.raises(ValidationError, match="embedding_model must not be blank"): cfg.embedding_model = "\t" + def test_fusion_config_fields_enforce_their_bounds(self): + """The new fusion/expansion knobs reject out-of-range values, so a bad + config surfaces at assignment instead of silently mis-weighting fusion.""" + from pydantic import ValidationError + + for field, bad in [ + ("lexical_fusion_weight", 1.5), + ("lexical_fusion_weight", -0.1), + ("adaptive_fusion_margin", 2.5), + ("adaptive_fusion_margin", -0.1), + ("title_search_weight", 1.5), + ("title_search_weight", -0.1), + ("neighbor_expansion", -1), + ("neighbor_expansion", 101), # upper bound guards against a token-count misread + ]: + with pytest.raises(ValidationError): + setattr(cfg, field, bad) + def test_embedding_dim_override(self): with mock.patch.dict(os.environ, {"LILBEE_EMBEDDING_DIM": "1024"}): c = Config() diff --git a/tests/test_store.py b/tests/test_store.py index baa1d088d..6057547e4 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -437,6 +437,64 @@ def test_has_scalar_index_is_false_when_listing_raises(self, store): with mock.patch.object(type(table), "list_indices", side_effect=RuntimeError("boom")): assert _has_scalar_index(table, "source") is False + def test_indexes_chunk_concepts_source_column(self, store): + """The concept-boost path filters chunk_concepts by chunk_source per + result, so that column gets its own BTree index too.""" + from lilbee.core.config import CHUNK_CONCEPTS_TABLE + from lilbee.data.store import ensure_table + from lilbee.data.store.lance_helpers import _has_scalar_index + from lilbee.retrieval.concepts.schema import _chunk_concepts_schema + + store.add_chunks(_make_records()) + cc = ensure_table(store.get_db(), CHUNK_CONCEPTS_TABLE, _chunk_concepts_schema()) + cc.add([{"chunk_source": "doc.md", "chunk_index": 0, "concept": "alpha"}]) + store.ensure_scalar_indexes() + assert _has_scalar_index(store.open_table(CHUNK_CONCEPTS_TABLE), "chunk_source") + + def test_one_column_failure_does_not_skip_the_other(self, store): + """A BTree failure on 'source' must not skip the independent Bitmap on + 'chunk_type' -- each column gets its own try.""" + store.add_chunks(_make_records()) + table = store.open_table("chunks") + attempted = [] + + def _record(self, column, **kwargs): + attempted.append(column) + if column == "source": + raise RuntimeError("boom") + + with mock.patch.object(type(table), "create_scalar_index", _record): + store.ensure_scalar_indexes() + assert attempted == ["source", "chunk_type"] + + def test_scalar_index_failure_on_populated_table_warns(self, store, caplog): + """A create failure on a non-empty table warns (it silently loses the + prefilter speedup), unlike the benign empty-table case.""" + import logging + + store.add_chunks(_make_records()) + table = store.open_table("chunks") + with ( + mock.patch.object(type(table), "create_scalar_index", side_effect=RuntimeError("boom")), + caplog.at_level(logging.WARNING), + ): + store.ensure_scalar_indexes() + warns = [r for r in caplog.records if r.levelno >= logging.WARNING] + assert any("Scalar index create failed" in r.message for r in warns) + + def test_search_builds_scalar_indexes_on_a_serve_only_store(self, store, test_config): + """A store served without a fresh ingest never ran the ingest path that + builds scalar indexes, so the first search must build them once.""" + store.add_chunks(_make_records()) + assert store._scalar_ready is False # ingest path did not run here + with mock.patch.object( + store, "ensure_scalar_indexes", wraps=store.ensure_scalar_indexes + ) as spy: + store.search([0.5] * test_config.embedding_dim, top_k=3) + store.search([0.5] * test_config.embedding_dim, top_k=3) + spy.assert_called_once() # built once, then the guard skips it + assert store._scalar_ready is True + class TestEnsureVectorIndex: """Small vaults stay on exact flat search; large ones get an ANN index.""" @@ -2657,8 +2715,12 @@ def test_bm25_probe_stays_chunk_scoped(self, store): assert store.bm25_probe("zebra") == [] assert store.bm25_probe("plain body words") - def test_title_index_failure_never_blocks_chunk_index(self, store, test_config): - """A failing title index leaves chunk FTS ready; the arm degrades to empty.""" + def test_title_index_failure_never_blocks_chunk_index(self, store, test_config, caplog): + """A failing title index leaves chunk FTS ready; the arm degrades to + empty. Because the arm is enabled, the failure warns (not debug) so an + opted-in title arm that cannot build is not a silent no-op.""" + import logging + test_config.title_search = True store.add_chunks(_titled_records("a.pdf", 1, title="zebra manifesto")) table = store.open_table("chunks") @@ -2669,10 +2731,15 @@ def _fail_title(self, column, **kwargs): raise RuntimeError("boom") return real_create(self, column, **kwargs) - with mock.patch.object(type(table), "create_fts_index", _fail_title): + with ( + mock.patch.object(type(table), "create_fts_index", _fail_title), + caplog.at_level(logging.WARNING), + ): store.ensure_fts_index() assert store._fts_ready assert store._title_arm(store.open_table("chunks"), "zebra", 5, None) == [] + warnings = [r for r in caplog.records if r.levelno >= logging.WARNING] + assert any("title" in r.message.lower() for r in warnings) class TestSourceMetadata: From 9b3ee50193983b57bc1d3c2093416276d53dcccc Mon Sep 17 00:00:00 2001 From: Tobias Perelstein <5562156+tobocop2@users.noreply.github.com> Date: Sun, 19 Jul 2026 23:17:37 -0400 Subject: [PATCH 19/28] Collapse the title arm to one deterministic row per matched document The pipeline stamps a document's title onto every chunk, so all of its chunks tie on the title BM25 score. _title_arm's plain limit therefore returned an arbitrary, implementation-defined subset of one document's chunks, which fusion then ranked 1..k and which (carrying bm25_score) bypassed the max_distance cutoff -- the pure "query names a document by title" case the feature exists for. Per-chunk stamping also depressed large documents' title IDF, biasing them below small documents. The arm now over-fetches, collapses each source to one deterministic representative (its first chunk), and returns the top documents by title score. A title match surfaces each document once, stably, and large and small documents compete as one row each. Tests cover the single- and multi-document cases and determinism across calls. Addresses bb-s1x8. make check pending; targeted suites green. --- src/lilbee/data/store/core.py | 30 ++++++++++++++++++++++++++++-- tests/test_store.py | 32 ++++++++++++++++++++++++++++++-- 2 files changed, 58 insertions(+), 4 deletions(-) diff --git a/src/lilbee/data/store/core.py b/src/lilbee/data/store/core.py index 8488c72d9..6cc71d0a1 100644 --- a/src/lilbee/data/store/core.py +++ b/src/lilbee/data/store/core.py @@ -182,6 +182,12 @@ def _lexical_rows( # bounds the decoded-text working set while the scan stays columnar. _TERM_SCAN_BATCH_ROWS = 20_000 +# The title arm collapses each matched document to one row, so it over-fetches +# to gather enough distinct documents before deduping. Bounded so a title that +# hits a huge document can't scan the whole corpus. +_TITLE_FETCH_FACTOR = 20 +_TITLE_MIN_FETCH = 200 + def _ann_nprobes(row_count: int) -> int: """Partitions to probe: a fixed fraction of the IVF partition count (~sqrt(N)), floored.""" @@ -792,7 +798,14 @@ def _title_arm( limit: int, chunk_type: ChunkType | None, ) -> list[SearchChunk]: - """BM25-arm candidates over the document title column. + """One BM25 row per document whose title matches, in title-relevance order. + + Every chunk of a document carries the same title, so all of its chunks + tie on BM25 and a plain ``limit`` would return an arbitrary tie-ordered + subset of a single document. Instead this over-fetches, collapses each + source to one deterministic representative (its first chunk), and returns + the top *limit* documents ordered by title score -- so "a query naming a + document by title surfaces its chunks" holds as one stable row per doc. Empty when the store predates the title column or its FTS index (old indexes keep working) and empty on any query-time failure: the optional @@ -802,10 +815,23 @@ def _title_arm( if not _has_fts_index(table, _TITLE_COLUMN): return [] try: - return _lexical_rows(table, query_text, limit, chunk_type, column=_TITLE_COLUMN) + rows = _lexical_rows( + table, + query_text, + max(limit * _TITLE_FETCH_FACTOR, _TITLE_MIN_FETCH), + chunk_type, + column=_TITLE_COLUMN, + ) except Exception: log.debug("Title arm search failed; contributing no title rows", exc_info=True) return [] + best: dict[str, SearchChunk] = {} + for row in rows: + seen = best.get(row.source) + if seen is None or row.chunk_index < seen.chunk_index: + best[row.source] = row + ordered = sorted(best.values(), key=lambda r: (-(r.bm25_score or 0.0), r.source)) + return ordered[:limit] def _hybrid_search( self, diff --git a/tests/test_store.py b/tests/test_store.py index 6057547e4..6cfe88941 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -2589,7 +2589,9 @@ class TestTitleSearch: def test_title_arm_surfaces_title_only_match(self, store, test_config): """A term that lives only in a document's title reaches hybrid results - with lexical support (bm25_score) when title_search is on.""" + with lexical support (bm25_score) when title_search is on. The title arm + surfaces one representative chunk per matched document, so the document + appears with lexical support even though not every chunk carries it.""" test_config.title_search = True store.add_chunks(_titled_records("a.pdf", 2, title="zebra manifesto", base=0.9)) store.add_chunks(_titled_records("b.pdf", 2, title="meeting notes", base=0.1)) @@ -2598,7 +2600,33 @@ def test_title_arm_surfaces_title_only_match(self, store, test_config): results = store.search(query_vec, top_k=4, max_distance=0, query_text="zebra") matched = [r for r in results if r.source == "a.pdf"] assert matched - assert all(r.bm25_score is not None for r in matched) + assert any(r.bm25_score is not None for r in matched) + + def test_title_arm_collapses_to_one_deterministic_row_per_document(self, store, test_config): + """Every chunk of a document shares its title, so all tie on BM25. The + arm must collapse each matched document to one deterministic row (its + first chunk), not return an arbitrary tie-ordered subset of that doc.""" + test_config.title_search = True + store.add_chunks(_titled_records("a.pdf", 5, title="zebra manifesto")) + store.ensure_fts_index() + table = store.open_table("chunks") + rows = store._title_arm(table, "zebra", 5, None) + assert [(r.source, r.chunk_index) for r in rows] == [("a.pdf", 0)] + # Deterministic across repeated calls (no implementation-defined tie order). + again = store._title_arm(table, "zebra", 5, None) + assert [(r.source, r.chunk_index) for r in again] == [("a.pdf", 0)] + + def test_title_arm_one_row_per_matched_document(self, store, test_config): + """Two matched documents surface one representative row each, not an + arbitrary flood of one document's chunks.""" + test_config.title_search = True + store.add_chunks(_titled_records("zebra.pdf", 4, title="zebra")) + store.add_chunks(_titled_records("safari.pdf", 4, title="zebra safari")) + store.ensure_fts_index() + table = store.open_table("chunks") + rows = store._title_arm(table, "zebra", 5, None) + assert sorted(r.source for r in rows) == ["safari.pdf", "zebra.pdf"] + assert all(r.chunk_index == 0 for r in rows) def test_title_search_weight_reaches_fusion(self, store, test_config): """A non-default title_search_weight is threaded into fuse_arms, not From 4949c2cbb5d4cae919c31981999884c3bdc46a2e Mon Sep 17 00:00:00 2001 From: Tobias Perelstein <5562156+tobocop2@users.noreply.github.com> Date: Sun, 19 Jul 2026 23:51:37 -0400 Subject: [PATCH 20/28] Round-trip source metadata through export/import Export wrote the per-page dataset with no source metadata, so import could only stamp the filename stem and wrote empty authors/created_at. A vault whose documents had real extracted titles lost them on an export/import cycle, and re-importing over an existing source overwrote its populated columns with stem/NULL/NULL. The dataset now carries title/authors/created_at denormalized onto each page row (read from the sources table), and import restores them, falling back to the stem-derived title when a dataset predates the columns. The page-texts write projects the row back down to that table's four columns. This also gives the _sources title/authors/created_at columns their first real reader: export consumes them, so they are no longer write-only schema. Addresses bb-vofd and bb-r1c3. make check pending; targeted suites green. --- src/lilbee/data/export.py | 61 +++++++++++++++++++++++++++++++--- src/lilbee/data/store/types.py | 11 +++++- tests/test_export.py | 39 ++++++++++++++++++++-- 3 files changed, 103 insertions(+), 8 deletions(-) diff --git a/src/lilbee/data/export.py b/src/lilbee/data/export.py index ab541960a..3754ee815 100644 --- a/src/lilbee/data/export.py +++ b/src/lilbee/data/export.py @@ -95,9 +95,27 @@ def build_page_dataset(store: Store, source: str | None = None) -> pa.Table: extra = _reconstructed_arrow(store, sorted(names - captured), table.schema) if extra.num_rows: table = pa.concat_tables([table, extra]) + table = _with_source_metadata(store, table) return table.sort_by([("source", "ascending"), ("page", "ascending")]) +def _with_source_metadata(store: Store, table: pa.Table) -> pa.Table: + """Denormalize each source's extraction metadata onto its page rows. + + The dataset is per page while title/authors/created_at live per source, so + without this an export/import cycle drops them and the import can only fall + back to the filename stem. Absent metadata stays null. + """ + import pyarrow as pa + + by_name: dict[str, dict] = {s["filename"]: dict(s) for s in store.get_sources()} + sources = table.column("source").to_pylist() + for column in SourceMeta._fields: + values = [by_name.get(name, {}).get(column) for name in sources] + table = table.append_column(column, pa.array(values, pa.string())) + return table + + def _reconstructed_arrow(store: Store, sources: list[str], schema: pa.Schema) -> pa.Table: """Chunk-reconstructed pages for *sources* as an Arrow table in *schema*.""" import pyarrow as pa @@ -201,6 +219,37 @@ def load_page_dataset(path: Path, fmt: DatasetFormat) -> list[PageTextRecord]: return deserialize_dataset(path.read_bytes(), fmt) +def _page_text_row(row: PageTextRecord) -> dict: + """Project a dataset row down to the ``_page_texts`` columns. + + A dataset carries the source's metadata denormalized on every page row; the + page-texts table has no such columns, so they are dropped before the write. + """ + return { + "source": row["source"], + "page": row["page"], + "text": row["text"], + "content_type": row["content_type"], + } + + +def _source_meta_from_rows(rows: list[PageTextRecord], name: str) -> SourceMeta: + """Recover a source's extraction metadata from its dataset rows. + + The values are identical on every page row, so the first carries them. A + dataset exported before the metadata columns existed has none, in which case + the title falls back to the cleaned filename stem. + """ + first: dict = dict(rows[0]) if rows else {} + stored = first.get("title") + title = stored.strip() if isinstance(stored, str) and stored.strip() else derive_title(name) + return SourceMeta( + title=title, + authors=first.get("authors") or "", + created_at=first.get("created_at") or "", + ) + + async def import_dataset( store: Store, rows: list[PageTextRecord], @@ -226,9 +275,11 @@ async def import_dataset( content_type = source_rows[0]["content_type"] or "text" page_texts = [(r["page"], r["text"]) for r in source_rows] chunks = await chunk_and_embed_pages(page_texts, name, content_type, on_progress) - # Imports carry no extraction metadata; the stem-derived title keeps - # imported chunks visible to the title search arm. - title = derive_title(name) + # Datasets exported with the metadata columns round-trip the extracted + # title/authors/created_at; older ones carry none, so fall back to the + # stem-derived title that keeps imported chunks visible to the title arm. + meta = _source_meta_from_rows(source_rows, name) + title = meta.title for chunk in chunks: chunk["title"] = title or None # One locked transaction (cleanup + chunks + page texts + source row) so a @@ -242,9 +293,9 @@ async def import_dataset( file_hash="", records=cast(list[dict], chunks), needs_cleanup=True, - page_texts=[dict(r) for r in source_rows], + page_texts=[_page_text_row(r) for r in source_rows], source_type=SourceType.IMPORTED, - meta=SourceMeta(title=title), + meta=meta, ) ], ) diff --git a/src/lilbee/data/store/types.py b/src/lilbee/data/store/types.py index 2d380825f..8f0ab2a61 100644 --- a/src/lilbee/data/store/types.py +++ b/src/lilbee/data/store/types.py @@ -260,12 +260,21 @@ class SourceStatBackfill(NamedTuple): class PageTextRecord(TypedDict): - """One row of the per-page text dataset, matching ``_page_texts``.""" + """One row of the per-page text dataset, matching ``_page_texts``. + + The export dataset additionally carries the source's extraction metadata + (denormalized onto every page row) so an export/import cycle preserves it; + these are absent on rows read from the ``_page_texts`` table and on datasets + exported before the columns existed. + """ source: str page: int text: str content_type: str + title: NotRequired[str | None] + authors: NotRequired[str | None] + created_at: NotRequired[str | None] class CitationRecord(TypedDict): diff --git a/tests/test_export.py b/tests/test_export.py index a022796cd..4be5f1bbb 100644 --- a/tests/test_export.py +++ b/tests/test_export.py @@ -265,10 +265,45 @@ def counting_batch(items): assert store.get_sources()[0]["source_type"] == SourceType.IMPORTED async def test_import_stamps_stem_title(self, services): - # Imports carry no extraction metadata; the stem-derived title keeps - # imported chunks visible to the title search arm. + # A dataset exported before the metadata columns existed carries none, so + # the stem-derived title keeps imported chunks visible to the title arm. store = services await import_dataset(store, [_page("field_notes.pdf", 1, "page body")]) chunks = store.get_chunks_by_source("field_notes.pdf") assert chunks and all(c.title == "field notes" for c in chunks) assert store.get_sources()[0]["title"] == "field notes" + + async def test_import_restores_extraction_metadata_from_the_dataset(self, services): + """A dataset carrying the metadata columns restores the real extracted + title/authors/created_at instead of downgrading to the filename stem.""" + store = services + row = { + **_page("report_2021.pdf", 1, "page body"), + "title": "Annual Report", + "authors": "Ada, Grace", + "created_at": "2021-05-01", + } + await import_dataset(store, [row]) + chunks = store.get_chunks_by_source("report_2021.pdf") + assert chunks and all(c.title == "Annual Report" for c in chunks) + source = store.get_sources()[0] + assert source["title"] == "Annual Report" + assert source["authors"] == "Ada, Grace" + assert source["created_at"] == "2021-05-01" + + async def test_export_round_trips_the_metadata_back_out(self, services): + """The exported dataset carries each source's metadata on its page rows, + so an export/import cycle preserves it instead of losing it.""" + store = services + row = { + **_page("report_2021.pdf", 1, "page body"), + "title": "Annual Report", + "authors": "Ada, Grace", + "created_at": "2021-05-01", + } + await import_dataset(store, [row]) + exported = build_page_dataset(store).to_pylist() + assert exported + assert all(r["title"] == "Annual Report" for r in exported) + assert all(r["authors"] == "Ada, Grace" for r in exported) + assert all(r["created_at"] == "2021-05-01" for r in exported) From e1a723aaf8df4e87c361f7a00253b43ddafdef9e Mon Sep 17 00:00:00 2001 From: Tobias Perelstein <5562156+tobocop2@users.noreply.github.com> Date: Mon, 20 Jul 2026 00:16:13 -0400 Subject: [PATCH 21/28] Neighbor expansion: version-guard, linear overlap, single budget accounting Addresses the neighbor/context P3 cluster on feat/retrieval-consolidated: - bb-0hrq: neighbor fetch now reads the center rows too, and _widen skips widening when a center's stored text no longer matches the search snapshot (a concurrent re-ingest re-chunked the file), so a passage is never spliced across document versions. - bb-cuw4: expand_neighbors short-circuits on a zero budget before paying the per-source store fetches; the tightest windows stop doing the most expensive reads for a guaranteed no-op. - bb-inur: _overlap_chars is now a linear prefix-function computation instead of trying every suffix length, so the common no-overlap seam is O(n) not O(n^2). - bb-td06: _fit_to_budget returns the tokens spent so _finalize_context derives the neighbor leftover once, instead of _widen_with_neighbors re-deriving the same budget and re-summing the same per-chunk cost. - bb-q3on: _widen_with_neighbors docstring no longer claims the sources block is untouched (the branch's own test shows the widened page span renders). - bb-jpcf: dropped the redundant per-field cfg save/restore fixture (conftest's _isolate_cfg snapshots the whole model). - bb-odgx: the optimize()-ordering comment now states the structural invariant rather than an ordering the inner try/except already carries. - bb-v2to: raw "chunk" FTS-column literals replaced with a _CHUNK_COLUMN constant beside _TITLE_COLUMN. make check pending; targeted suites green (550 passed). --- src/lilbee/data/store/core.py | 14 ++++---- src/lilbee/data/store/lance_helpers.py | 7 +++- src/lilbee/retrieval/query/neighbors.py | 41 +++++++++++++++++---- src/lilbee/retrieval/query/searcher.py | 48 +++++++++++++------------ tests/test_neighbors.py | 37 ++++++++++++++++--- tests/test_query.py | 27 +++++--------- 6 files changed, 114 insertions(+), 60 deletions(-) diff --git a/src/lilbee/data/store/core.py b/src/lilbee/data/store/core.py index 6cc71d0a1..816e65226 100644 --- a/src/lilbee/data/store/core.py +++ b/src/lilbee/data/store/core.py @@ -30,6 +30,7 @@ from .fusion import adaptive_weight_scale, fuse_arms, normalized_bm25, vector_similarity from .lance_helpers import ( + _CHUNK_COLUMN, _chunk_type_predicate, _has_fts_index, _has_scalar_index, @@ -125,7 +126,7 @@ def _lexical_rows( query_text: str, limit: int, chunk_type: ChunkType | None, - column: str = "chunk", + column: str = _CHUNK_COLUMN, ) -> list[SearchChunk]: """BM25 rows for *query_text* over a single FTS *column*. @@ -480,10 +481,9 @@ def ensure_fts_index(self) -> None: return try: if _has_fts_index(table): - # Mark hybrid ready BEFORE optimize(): the existing index - # already serves queries, and optimize() can raise on a large - # corpus. A failure here must not drop every query to - # vector-only. + # An existing index serves queries regardless of how the + # best-effort optimize() below turns out, so hybrid is ready + # either way. self._fts_ready = True try: # One optimize folds new rows into every index on the @@ -507,7 +507,7 @@ def ensure_fts_index(self) -> None: # encoding on optimize() for a large corpus. Positions only # serve exact-phrase queries, which lilbee never issues (FTS # queries match plain terms), so the index never needs them. - table.create_fts_index("chunk", replace=False, with_position=False) + table.create_fts_index(_CHUNK_COLUMN, replace=False, with_position=False) self._fts_ready = True log.debug("FTS index created on '%s'", CHUNKS_TABLE) # Only the opt-in title arm needs the title index; without this a @@ -547,7 +547,7 @@ def _rebuild_fts_positionless(self, table: lancedb.table.Table) -> None: title index is rebuilt too when the title arm is enabled. """ try: - table.create_fts_index("chunk", replace=True, with_position=False) + table.create_fts_index(_CHUNK_COLUMN, replace=True, with_position=False) if self._config.title_search and _TITLE_COLUMN in table.schema.names: table.create_fts_index(_TITLE_COLUMN, replace=True, with_position=False) log.warning("Rebuilt the FTS index positionless after a positional-index overflow") diff --git a/src/lilbee/data/store/lance_helpers.py b/src/lilbee/data/store/lance_helpers.py index 265e3ace3..6cfdda61e 100644 --- a/src/lilbee/data/store/lance_helpers.py +++ b/src/lilbee/data/store/lance_helpers.py @@ -20,6 +20,11 @@ log = logging.getLogger(__name__) +# The chunks table's body-text column: FTS-indexed and searched by the lexical +# arm, named here so the store's query and index code shares one spelling with +# its sibling _TITLE_COLUMN. +_CHUNK_COLUMN = "chunk" + def install_lancedb_thread_error_suppressor() -> None: """Install a ``threading.excepthook`` that swallows lancedb shutdown noise. @@ -128,7 +133,7 @@ def _chunk_type_predicate(chunk_type: ChunkType | str) -> str: return f"chunk_type = '{escaped}'" -def _has_fts_index(table: lancedb.table.Table, column: str = "chunk") -> bool: +def _has_fts_index(table: lancedb.table.Table, column: str = _CHUNK_COLUMN) -> bool: """Return True when an FTS index on *column* already exists.""" try: for idx in table.list_indices(): diff --git a/src/lilbee/retrieval/query/neighbors.py b/src/lilbee/retrieval/query/neighbors.py index e874b6785..11eade2be 100644 --- a/src/lilbee/retrieval/query/neighbors.py +++ b/src/lilbee/retrieval/query/neighbors.py @@ -25,11 +25,25 @@ def _overlap_chars(left: str, right: str) -> int: 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. + + Computed with a prefix-function pass rather than trying every length: the + no-overlap case (a ``chunk_overlap=0`` build) is common and would otherwise + scan every length to exhaustion, quadratic in the chunk's own size. """ - for k in range(min(len(left), len(right)), 0, -1): - if left.endswith(right[:k]): - return k - return 0 + if not left or not right: + return 0 + # Only left's last len(right) chars can match a prefix of right, so the + # probe stays linear in the incoming chunk, not the accumulated passage. + probe = f"{right}\0{left[-len(right) :]}" + failure = [0] * len(probe) + length = 0 + for i in range(1, len(probe)): + while length and probe[i] != probe[length]: + length = failure[length - 1] + if probe[i] == probe[length]: + length += 1 + failure[i] = length + return failure[-1] def merge_adjacent_texts(texts: list[str]) -> str: @@ -65,6 +79,11 @@ def expand_neighbors( higher-ranked expansion, is never pulled again, so no passage text is duplicated (a document routed whole expands to nothing). """ + if budget <= 0: + # Nothing can be spent, so skip the per-source store fetches entirely: + # on a tight window every widen attempt would fail against a zero + # budget after paying for the reads. + return results centers: dict[str, set[int]] = {} for r in results: centers.setdefault(r.source, set()).add(r.chunk_index) @@ -86,8 +105,9 @@ def _fetch_neighbor_rows( ) -> 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. + The centers are fetched alongside their neighbors so the caller can tell + whether the document was re-ingested since the search snapshot; 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(): @@ -96,7 +116,7 @@ def _fetch_neighbor_rows( index for center in owned for index in range(center - radius, center + radius + 1) - if index >= 0 and index not in owned + if index >= 0 } ) for row in store.get_chunks_by_indices(source, wanted): @@ -138,6 +158,13 @@ def _widen( original chunk untouched. """ center = result.chunk_index + current = rows.get((result.source, center)) + if current is not None and current.chunk != result.chunk: + # The document was re-ingested between the search and this fetch (reads + # take no lock and the store's read-consistency window is seconds), so + # the neighbor rows belong to a different chunking of the file. Splicing + # them would invent text and a page span that existed in no version. + return result, 0 left = _neighbor_run(result, rows, claimed, -1, radius) right = _neighbor_run(result, rows, claimed, +1, radius) while left or right: diff --git a/src/lilbee/retrieval/query/searcher.py b/src/lilbee/retrieval/query/searcher.py index 234f0a509..f004b0329 100644 --- a/src/lilbee/retrieval/query/searcher.py +++ b/src/lilbee/retrieval/query/searcher.py @@ -938,8 +938,9 @@ def _finalize_context( """ system = self._system_with_memory(self._config.rag_system_prompt, question) base_results = list(results) - results = self._fit_context_budget(results, system, question, history, scale) - results = self._widen_with_neighbors(results, system, question, history, scale) + budget = self._context_budget(system, question, history, scale) + results, used = self._fit_to_budget(results, budget) + results = self._widen_with_neighbors(results, max(0, budget - used)) context = build_context(results) prompt = CONTEXT_TEMPLATE.format(context=context, question=question) messages: list[ChatMessage] = [{"role": "system", "content": system}] @@ -996,7 +997,19 @@ def _fit_context_budget( 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) + return self._fit_to_budget(results, self._context_budget(system, question, history, scale))[ + 0 + ] + + def _fit_to_budget( + self, results: list[SearchChunk], budget: int + ) -> tuple[list[SearchChunk], int]: + """Fit *results* into *budget*: the kept sources and the tokens they cost. + + Returning the spent total lets the caller derive the leftover for + neighbor expansion instead of re-deriving the same per-chunk cost, so + the two stages cannot drift apart on the accounting. + """ kept: list[SearchChunk] = [] used = 0 for r in results: @@ -1011,32 +1024,21 @@ def _fit_context_budget( len(kept), len(results), ) - return kept + return kept, used - def _widen_with_neighbors( - self, - results: list[SearchChunk], - system: str, - question: str, - history: list[ChatMessage] | None, - scale: float, - ) -> list[SearchChunk]: + def _widen_with_neighbors(self, results: list[SearchChunk], leftover: int) -> 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. + Spends only *leftover*, the budget the fit did not use, so a tight + window sheds expansion first and never drops an original chunk for a + neighbor. Widening keeps each passage's citation number and identity; + its text and page/line span do change, so the sources block shows the + widened range. """ radius = self._config.neighbor_expansion - if radius <= 0: + if radius <= 0 or leftover <= 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 - ) + return expand_neighbors(results, self._store, radius, leftover, 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.""" diff --git a/tests/test_neighbors.py b/tests/test_neighbors.py index c172cfcdc..63c103152 100644 --- a/tests/test_neighbors.py +++ b/tests/test_neighbors.py @@ -79,7 +79,7 @@ def test_widens_both_sides_and_recomputes_the_page_span(self): 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]) + store.get_chunks_by_indices.assert_called_once_with("doc.pdf", [1, 2, 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) @@ -109,7 +109,7 @@ def test_selected_neighbors_are_never_pulled_twice(self): 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]) + store.get_chunks_by_indices.assert_called_once_with("doc.pdf", [1, 2, 3, 4]) def test_higher_ranked_expansion_claims_the_shared_neighbor(self): first = _chunk(index=2, text="bb") @@ -135,9 +135,38 @@ def test_budget_sheds_farthest_neighbors_first_trailing_on_ties(self): 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) + store = _store_with(rows) + out = expand_neighbors([center], store, radius=1, budget=0, cost=_cost) + assert out == [center] + # A zero budget can buy nothing, so it must not pay for the store reads. + store.get_chunks_by_indices.assert_not_called() + + def test_reingested_document_is_not_spliced_across_versions(self): + """If the file is re-ingested between the search and the neighbor fetch, + the rows describe a different chunking. Splicing them would invent text + and a page span that existed in no version, so widening is skipped.""" + center = _chunk(index=2, text="original center text") + rows = [ + _chunk(index=1, text="new left"), + # Same (source, chunk_index), different text: the file was re-chunked. + _chunk(index=2, text="re-ingested center text"), + _chunk(index=3, text="new right"), + ] + out = expand_neighbors([center], _store_with(rows), radius=1, budget=1000, cost=_cost) assert out == [center] + def test_unchanged_center_still_widens(self): + """The version guard must not block the normal path: a center whose + stored row still matches widens as before.""" + center = _chunk(index=2, text="beta gamma") + rows = [ + _chunk(index=1, text="alpha beta"), + _chunk(index=2, text="beta gamma"), + _chunk(index=3, text="gamma delta"), + ] + out = expand_neighbors([center], _store_with(rows), radius=1, budget=1000, cost=_cost) + assert out[0].chunk != "beta gamma" + 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. @@ -145,7 +174,7 @@ def test_whole_document_selection_expands_nothing(self): 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]) + store.get_chunks_by_indices.assert_called_once_with("doc.pdf", [0, 1, 2, 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, diff --git a/tests/test_query.py b/tests/test_query.py index 6ba790ed8..c8155041d 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -521,18 +521,14 @@ def test_structural_filter_runs_before_the_buffer_slice_so_real_passages_backfil """The filter drops structural chunks from the candidate list before the top_k*2 slice, so a real passage past the window backfills into it. If the filter ran after the slice, the deeper real passage would be lost.""" - old_top_k = cfg.top_k cfg.top_k = 1 # window is top_k*2 = 2 cfg.filter_structural_chunks = True - try: - real0 = _make_result(source="a.pdf", distance=0.1, chunk="Answer prose one.") - toc1 = _make_result(source="toc1.pdf", distance=0.2, chunk=self._TOC_CHUNK) - toc2 = _make_result(source="toc2.pdf", distance=0.3, chunk=self._TOC_CHUNK) - real3 = _make_result(source="b.pdf", distance=0.4, chunk="Answer prose two.") - mock_svc.store.search.return_value = [real0, toc1, toc2, real3] - sources = [r.source for r in get_services().searcher.search("q")] - finally: - cfg.top_k = old_top_k + real0 = _make_result(source="a.pdf", distance=0.1, chunk="Answer prose one.") + toc1 = _make_result(source="toc1.pdf", distance=0.2, chunk=self._TOC_CHUNK) + toc2 = _make_result(source="toc2.pdf", distance=0.3, chunk=self._TOC_CHUNK) + real3 = _make_result(source="b.pdf", distance=0.4, chunk="Answer prose two.") + mock_svc.store.search.return_value = [real0, toc1, toc2, real3] + sources = [r.source for r in get_services().searcher.search("q")] assert "b.pdf" in sources # backfilled from past the pre-filter window assert "toc1.pdf" not in sources and "toc2.pdf" not in sources @@ -951,12 +947,6 @@ def test_logs_when_trimming(self, mock_svc, caplog): 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") @@ -1064,8 +1054,9 @@ def test_widened_prompt_still_fits_the_provider_ceiling(self, mock_svc): base = [_make_result(chunk="c" * 1500, chunk_index=2, page_start=3, page_end=3)] searcher = get_services().searcher - fitted = searcher._fit_context_budget(base, system, question, None, 1.0) - widened = searcher._widen_with_neighbors(fitted, system, question, None, 1.0) + budget = searcher._context_budget(system, question, None, 1.0) + fitted, used = searcher._fit_to_budget(base, budget) + widened = searcher._widen_with_neighbors(fitted, max(0, budget - used)) # Non-vacuous: at least one neighbor was actually merged in. assert widened[0].chunk != "c" * 1500 From a9471998eb2597adc29790a9d0e35395fdfcd427 Mon Sep 17 00:00:00 2001 From: Tobias Perelstein <5562156+tobocop2@users.noreply.github.com> Date: Mon, 20 Jul 2026 00:58:15 -0400 Subject: [PATCH 22/28] Ingest returns (records, metadata); markdown H1 and image EXIF titles Replaces the meta_out list out-parameter with a structured (records, SourceMeta) return through ingest_document, ingest_markdown, _handle_image, _produce_records, and ingest_batch, removing the positional [0] reads on single-element lists. With the plumbing carrying real metadata, two title gaps close: - Markdown notes now take their title from the leading `# Heading` (the best title a note carries), falling back to the filename stem when there is none. This is the flagship .md corpus's highest-quality title source. - Images capture their EXIF/XMP title/authors/date via a metadata-only kreuzberg read (the OCR path never touched kreuzberg), instead of always downgrading to the stem. Chunk titles still persist NULL for an empty title. Test doubles and callers updated to the tuple return; the markdown-title test now pins H1 extraction. Addresses bb-cp4z, bb-vv8k, bb-0rrl (and bb-wt45). make check pending; targeted suites green (609 passed). --- src/lilbee/data/ingest/extract.py | 81 ++++++++++++++++------ src/lilbee/data/ingest/pipeline.py | 33 ++++----- tests/test_ingest.py | 107 +++++++++++++++++++---------- tests/test_neighbors.py | 6 ++ tests/test_settings.py | 21 +----- tests/test_store.py | 5 +- 6 files changed, 152 insertions(+), 101 deletions(-) diff --git a/src/lilbee/data/ingest/extract.py b/src/lilbee/data/ingest/extract.py index 4f56fdabb..ca59f935a 100644 --- a/src/lilbee/data/ingest/extract.py +++ b/src/lilbee/data/ingest/extract.py @@ -22,7 +22,7 @@ from lilbee.data.ingest.discovery import file_hash from lilbee.data.ingest.ocr_cache import load_ocr_pages, ocr_cache_key, store_ocr_pages from lilbee.data.ingest.offload import to_ingest_thread -from lilbee.data.ingest.title import source_meta_from_extraction +from lilbee.data.ingest.title import derive_title, source_meta_from_extraction from lilbee.data.ingest.types import ( IMAGE_CONTENT_TYPE, MARKDOWN_OUTPUT, @@ -508,6 +508,22 @@ async def _handle_scanned_pdf_fallback( return chunks +def _image_meta(path: Path, source_name: str) -> SourceMeta: + """Extraction metadata (EXIF/XMP title, authors, date) for an image. + + The OCR path never touches kreuzberg, so this is a separate metadata-only + read; any failure degrades to the stem-derived title. + """ + try: + from kreuzberg import extract_file_sync + + result = extract_file_sync(str(path), config=extraction_config(ExtractMode.MARKDOWN)) + return source_meta_from_extraction(result.metadata or {}, source_name) + except Exception: + log.debug("Image metadata extraction failed for %s; using the stem title", source_name) + return SourceMeta(title=derive_title(source_name)) + + async def _handle_image( path: Path, source_name: str, @@ -515,23 +531,27 @@ async def _handle_image( *, on_progress: DetailedProgressCallback, page_texts_out: list[PageTextRecord] | None = None, -) -> list[ChunkRecord]: +) -> tuple[list[ChunkRecord], SourceMeta]: """OCR an image: vision OCR when a vision model is configured, else Tesseract. An image has no text layer to extract first, so it routes straight to OCR -- the same downstream call a PDF page hits after it is rasterized to an image. + Its EXIF/XMP metadata is captured separately so the title/authors/date are + not lost to the stem fallback. """ + meta = await to_ingest_thread(_image_meta, path, source_name) if _effective_enable_ocr() is False: # OCR explicitly disabled: an image has no text layer, so skip it rather # than paying the full Tesseract cost the config says is turned off. log.info("OCR disabled; skipping image OCR for %s", source_name) - return [] + return [], meta vision_model = active_config().vision_model if _should_run_ocr() and vision_model: log.info("Image: using vision OCR for %s (model=%s)", source_name, vision_model) - return await _vision_image_ocr( + chunks = await _vision_image_ocr( path, source_name, content_type, on_progress=on_progress, page_texts_out=page_texts_out ) + return chunks, meta log.info("Image: falling back to Tesseract OCR for %s", source_name) chunks = await _tesseract_ocr_fallback( @@ -539,7 +559,7 @@ async def _handle_image( ) if not chunks: _warn_empty_ocr(source_name, "images") - return chunks + return chunks, meta async def ingest_document( @@ -550,14 +570,12 @@ async def ingest_document( quiet: bool = False, on_progress: DetailedProgressCallback = noop_callback, page_texts_out: list[PageTextRecord] | None = None, - meta_out: list[SourceMeta] | None = None, -) -> list[ChunkRecord]: - """Extract and chunk a document, embed, return records. +) -> tuple[list[ChunkRecord], SourceMeta]: + """Extract and chunk a document, embed, and return (records, metadata). Vision OCR is controlled by ``cfg.enable_ocr`` (see ``_should_run_ocr``). - When ``page_texts_out`` is given, per-page text is appended for export. - When ``meta_out`` is given, the document's extraction metadata (title, - authors, creation date) is appended for the source row. + When ``page_texts_out`` is given, per-page text is appended for export. The + returned metadata carries the document's extraction title/authors/date. """ # An image carries no text layer; route it straight to OCR (vision or Tesseract) # instead of a no-op kreuzberg markdown extract that yields nothing for a scan. @@ -573,11 +591,10 @@ async def ingest_document( # Captured before the scanned-PDF fallback: a scan's PDF metadata (title, # authors) survives even when its text layer is empty. - if meta_out is not None: - meta_out.append(source_meta_from_extraction(result.metadata or {}, source_name)) + meta = source_meta_from_extraction(result.metadata or {}, source_name) if content_type == PDF_CONTENT_TYPE and not _has_meaningful_text(result): - return await _handle_scanned_pdf_fallback( + records = await _handle_scanned_pdf_fallback( path, source_name, content_type, @@ -586,9 +603,10 @@ async def ingest_document( on_progress=on_progress, page_texts_out=page_texts_out, ) + return records, meta if not result.chunks: - return [] + return [], meta _capture_result_page_texts(result, source_name, content_type, page_texts_out) @@ -608,7 +626,7 @@ async def ingest_document( get_services().embedder.embed_batch, texts, source=source_name, on_progress=on_progress ) - return [ + records = [ ChunkRecord( source=source_name, content_type=content_type, @@ -623,6 +641,23 @@ async def ingest_document( ) for idx, (chunk, text, vec) in enumerate(zip(result.chunks, texts, vectors, strict=True)) ] + return records, meta + + +def _markdown_h1(text: str) -> str | None: + """The document's leading ``# Heading``, the best title a note carries. + + Only a top-level ATX heading counts; ``##`` and deeper are sections, not the + document title. None when the note opens without one. + """ + for line in text.splitlines(): + stripped = line.strip() + if not stripped: + continue + if stripped.startswith("# ") and not stripped.startswith("## "): + return stripped[2:].strip() or None + return None + return None async def ingest_markdown( @@ -630,15 +665,18 @@ async def ingest_markdown( source_name: str, on_progress: DetailedProgressCallback = noop_callback, page_texts_out: list[PageTextRecord] | None = None, -) -> list[ChunkRecord]: +) -> tuple[list[ChunkRecord], SourceMeta]: """Chunk a markdown file with heading context prepended to each chunk. + Each chunk gets the heading hierarchy path (e.g. "# Setup > ## Install") prepended for better retrieval context. When ``page_texts_out`` is given, - the full text is appended as page 0 for export. + the full text is appended as page 0 for export. The returned metadata's + title is the note's leading ``# Heading`` when it has one, else the stem. """ raw_text = await to_ingest_thread(path.read_text, encoding="utf-8", errors="replace") + meta = SourceMeta(title=derive_title(source_name, _markdown_h1(raw_text))) if not raw_text.strip(): - return [] + return [], meta # chunk_text runs kreuzberg's synchronous extractor; offload it so a large # markdown doc does not stall sibling files sharing this event loop. @@ -646,7 +684,7 @@ async def ingest_markdown( chunk_text, raw_text, mime_type="text/markdown", heading_context=True ) if not texts: - return [] + return [], meta if page_texts_out is not None: page_texts_out.append(_page_text_record(source_name, 0, raw_text, "text")) @@ -654,7 +692,7 @@ async def ingest_markdown( vectors = await to_ingest_thread( get_services().embedder.embed_batch, texts, source=source_name, on_progress=on_progress ) - return [ + records = [ ChunkRecord( source=source_name, content_type="text", @@ -669,3 +707,4 @@ async def ingest_markdown( ) for idx, (t, vec) in enumerate(zip(texts, vectors, strict=True)) ] + return records, meta diff --git a/src/lilbee/data/ingest/pipeline.py b/src/lilbee/data/ingest/pipeline.py index 1daf2870b..b6035580f 100644 --- a/src/lilbee/data/ingest/pipeline.py +++ b/src/lilbee/data/ingest/pipeline.py @@ -200,45 +200,40 @@ async def _produce_records( quiet: bool = False, on_progress: DetailedProgressCallback = noop_callback, page_texts_out: list[PageTextRecord] | None = None, - meta_out: list[SourceMeta] | None = None, -) -> list[ChunkRecord]: - """Extract, chunk, and embed a single file into store-ready records. +) -> tuple[list[ChunkRecord], SourceMeta]: + """Extract, chunk, and embed a single file into (records, source metadata). The LanceDB write is deferred: records are returned to the caller and written in a batched flush (see :func:`_flush_writes`), so bulk ingest pays one write-lock acquisition per batch instead of one per file. The per-page text - dataset rows land in ``page_texts_out`` and are written by the same flush; - the document's metadata (extraction-provided when available, stem-derived - title otherwise) lands in ``meta_out`` and stamps every record's ``title``. + dataset rows land in ``page_texts_out`` and are written by the same flush. + The returned metadata (extraction-provided when available, stem-derived title + otherwise) stamps every record's ``title`` and updates the source row. """ records: list[ChunkRecord] page_texts: list[PageTextRecord] = page_texts_out if page_texts_out is not None else [] - meta = SourceMeta(title=derive_title(source_name)) if content_type == "code": records = await to_ingest_thread(ingest_code_sync, path, source_name, on_progress) + meta = SourceMeta(title=derive_title(source_name)) elif path.suffix.lower() == ".md": - records = await ingest_markdown(path, source_name, on_progress, page_texts_out=page_texts) + records, meta = await ingest_markdown( + path, source_name, on_progress, page_texts_out=page_texts + ) else: - extracted_meta: list[SourceMeta] = [] - records = await ingest_document( + records, meta = await ingest_document( path, source_name, content_type, quiet=quiet, on_progress=on_progress, page_texts_out=page_texts, - meta_out=extracted_meta, ) - if extracted_meta: - meta = extracted_meta[0] for record in records: # NULL (not "") for an absent title, so chunk rows match the migration # and the _sources table, which both persist absence as NULL. record["title"] = meta.title or None - if meta_out is not None: - meta_out.append(meta) - return records + return records, meta def _disk_stat(path: Path) -> SourceStat | None: @@ -697,15 +692,13 @@ async def _process_one(entry: FileToProcess, file_index: int) -> _IngestResult: # transaction as the new write (see _flush_writes), so cleanup is # carried on the result rather than run eagerly here. page_texts: list[PageTextRecord] = [] - meta_out: list[SourceMeta] = [] - records = await _produce_records( + records, meta = await _produce_records( entry.path, name, entry.content_type, quiet=quiet, on_progress=on_progress, page_texts_out=page_texts, - meta_out=meta_out, ) concept_records = await _build_concept_records(records, name) entity_rows = await _build_entity_records(records, name) @@ -726,7 +719,7 @@ async def _process_one(entry: FileToProcess, file_index: int) -> _IngestResult: stat=entry.stat, concept_records=concept_records, entity_rows=entity_rows, - meta=meta_out[0] if meta_out else None, + meta=meta, ) except (asyncio.CancelledError, TaskCancelledError) as exc: # TaskCancelledError is the TUI's cooperative cancel signal raised diff --git a/tests/test_ingest.py b/tests/test_ingest.py index 66043995e..3d28da7fb 100644 --- a/tests/test_ingest.py +++ b/tests/test_ingest.py @@ -897,7 +897,7 @@ async def testingest_document_empty_chunks(self, mock_extract_file, isolated_env f = isolated_env / "empty.txt" f.write_text(" ") - result = await ingest_document(f, "empty.txt", "text") + result, _ = await ingest_document(f, "empty.txt", "text") assert result == [] async def test_ingest_code_empty_chunks(self, isolated_env): @@ -956,7 +956,7 @@ async def testingest_document_pdf_with_pages(self, mock_kf, isolated_env): f = isolated_env / "test.pdf" f.write_bytes(b"fake") - result = await ingest_document(f, "test.pdf", "pdf") + result, _ = await ingest_document(f, "test.pdf", "pdf") assert len(result) == 2 assert result[0]["page_start"] == 1 assert result[1]["page_start"] == 2 @@ -1042,10 +1042,12 @@ class TestSkipMarkerLifecycle: until the file changes or retry_skipped / force_rebuild clears the marker.""" @staticmethod - def _zero_chunks(*_args, **_kwargs) -> list: + def _zero_chunks(*_args, **_kwargs): # Simulate "OCR found no usable text": no records produced, so the file # is recorded as skipped. - return [] + from lilbee.data.store import SourceMeta + + return [], SourceMeta() async def test_failed_file_is_skipped_on_next_sync(self, isolated_env, mock_svc): from lilbee.data.ingest import sync @@ -1099,13 +1101,14 @@ async def _pages_no_chunks( quiet=False, on_progress=None, page_texts_out=None, - meta_out=None, ): + from lilbee.data.store import SourceMeta + if page_texts_out is not None: page_texts_out.append( {"source": source_name, "page": 1, "text": " ", "content_type": "pdf"} ) - return [] + return [], SourceMeta() async def test_pages_and_source_row_persist_and_replan_stops(self, isolated_env, mock_svc): from lilbee.data.ingest import sync @@ -2196,7 +2199,7 @@ async def test_vision_fallback_called_for_empty_pdf(self, mock_kf, isolated_env, from lilbee.data.ingest import ingest_document - result = await ingest_document(f, "scanned.pdf", "pdf", quiet=True) + result, _ = await ingest_document(f, "scanned.pdf", "pdf", quiet=True) mock_svc.provider.pdf_ocr.assert_called_once_with( f, backend="vision", @@ -2246,7 +2249,7 @@ async def test_ocr_disabled_skips_both_backends(self, mock_kf, isolated_env, moc from lilbee.data.ingest import ingest_document with mock.patch("lilbee.data.ingest.extract._run_tesseract_sync") as mock_tess: - result = await ingest_document(f, "scanned.pdf", "pdf") + result, _ = await ingest_document(f, "scanned.pdf", "pdf") mock_svc.provider.pdf_ocr.assert_not_called() mock_tess.assert_not_called() # Only the initial text-layer extract ran; no Tesseract re-extract. @@ -2262,7 +2265,7 @@ async def test_vision_fallback_not_called_for_non_pdf(self, mock_kf, isolated_en from lilbee.data.ingest import ingest_document - result = await ingest_document(f, "doc.txt", "text") + result, _ = await ingest_document(f, "doc.txt", "text") mock_svc.provider.pdf_ocr.assert_not_called() assert result == [] @@ -2280,7 +2283,7 @@ async def test_vision_fallback_empty_vision_text_returns_empty( from lilbee.data.ingest import ingest_document - result = await ingest_document(f, "blank.pdf", "pdf") + result, _ = await ingest_document(f, "blank.pdf", "pdf") assert result == [] @mock.patch("kreuzberg.extract_file_sync", new_callable=Mock) @@ -2294,7 +2297,7 @@ async def test_no_vision_fallback_when_text_meaningful(self, mock_kf, isolated_e from lilbee.data.ingest import ingest_document - result = await ingest_document(f, "good.pdf", "pdf") + result, _ = await ingest_document(f, "good.pdf", "pdf") mock_svc.provider.pdf_ocr.assert_not_called() assert len(result) > 0 @@ -2311,7 +2314,7 @@ async def test_vision_fallback_no_chunks_returns_empty(self, mock_kf, isolated_e with mock.patch("lilbee.data.ingest.extract.chunk_text", return_value=[]): from lilbee.data.ingest import ingest_document - result = await ingest_document(f, "nochunks.pdf", "pdf") + result, _ = await ingest_document(f, "nochunks.pdf", "pdf") assert result == [] @mock.patch("kreuzberg.extract_file_sync", new_callable=Mock) @@ -2382,7 +2385,7 @@ async def test_image_routed_to_vision_ocr(self, isolated_env, mock_svc): from lilbee.data.ingest import ingest_document - result = await ingest_document(f, "scan.png", "image") + result, _ = await ingest_document(f, "scan.png", "image") # the single-image path is used, not the PDF page loop mock_svc.provider.vision_ocr.assert_called_once() mock_svc.provider.pdf_ocr.assert_not_called() @@ -2390,9 +2393,10 @@ async def test_image_routed_to_vision_ocr(self, isolated_env, mock_svc): assert result[0]["content_type"] == "image" assert result[0]["page_start"] == 1 - async def test_image_skips_kreuzberg_markdown_extract(self, isolated_env, mock_svc): - # The pre-fix bug: an image went through a markdown extract that yields no - # text. The image branch must not call kreuzberg.extract_file_sync at all. + async def test_image_text_comes_from_ocr_not_kreuzberg_extraction(self, isolated_env, mock_svc): + # The image's text must come from OCR, not a kreuzberg document extract + # (the pre-fix bug ran a markdown extract that yielded no text). kreuzberg + # is only used for a metadata-only read (title/authors from EXIF). cfg.vision_model = "org/Test-Vision-GGUF/test-vision-Q4_K_M.gguf" cfg.enable_ocr = True mock_svc.provider.vision_ocr.return_value = "text " * 20 @@ -2401,9 +2405,28 @@ async def test_image_skips_kreuzberg_markdown_extract(self, isolated_env, mock_s from lilbee.data.ingest import ingest_document - with mock.patch("kreuzberg.extract_file_sync", new_callable=Mock) as mock_kf: - await ingest_document(f, "scan.png", "image") - mock_kf.assert_not_called() + records, meta = await ingest_document(f, "scan.png", "image") + assert records # chunks were produced from OCR + assert mock_svc.provider.vision_ocr.called + assert meta.title == "scan" # no EXIF title in the test png -> stem + + async def test_image_metadata_failure_falls_back_to_the_stem_title( + self, isolated_env, mock_svc + ): + """A metadata read that raises must not fail the image: OCR still runs + and the title degrades to the filename stem.""" + cfg.vision_model = "org/Test-Vision-GGUF/test-vision-Q4_K_M.gguf" + cfg.enable_ocr = True + mock_svc.provider.vision_ocr.return_value = "text " * 20 + f = isolated_env / "broken_exif.png" + _write_png(f) + + from lilbee.data.ingest import ingest_document + + with mock.patch("kreuzberg.extract_file_sync", side_effect=RuntimeError("bad exif")): + records, meta = await ingest_document(f, "broken_exif.png", "image") + assert records + assert meta.title == "broken exif" async def test_image_falls_back_to_tesseract_without_vision_model(self, isolated_env, mock_svc): cfg.vision_model = "" @@ -2415,7 +2438,7 @@ async def test_image_falls_back_to_tesseract_without_vision_model(self, isolated with mock.patch("lilbee.data.ingest.extract._run_tesseract_sync") as mock_tess: mock_tess.return_value = _make_kreuzberg_result("Tesseract image text. " * 20) - result = await ingest_document(f, "scan.png", "image") + result, _ = await ingest_document(f, "scan.png", "image") mock_svc.provider.vision_ocr.assert_not_called() assert len(result) > 0 assert result[0]["content_type"] == "image" @@ -2448,7 +2471,7 @@ async def test_image_ocr_disabled_skips_both_backends(self, isolated_env, mock_s from lilbee.data.ingest import ingest_document with mock.patch("lilbee.data.ingest.extract._run_tesseract_sync") as mock_tess: - result = await ingest_document(f, "scan.png", "image") + result, _ = await ingest_document(f, "scan.png", "image") assert result == [] mock_svc.provider.vision_ocr.assert_not_called() mock_tess.assert_not_called() @@ -2489,7 +2512,8 @@ async def test_image_tesseract_empty_skips_file(self, isolated_env, mock_svc): with mock.patch("lilbee.data.ingest.extract._run_tesseract_sync") as mock_tess: mock_tess.return_value = _make_empty_result() - assert await ingest_document(f, "scan.png", "image") == [] + records, _ = await ingest_document(f, "scan.png", "image") + assert records == [] async def test_multipage_image_ocrs_every_frame_end_to_end(self, isolated_env, mock_svc): """A multi-frame TIFF routed to vision OCR yields one OCR call and one page @@ -2505,7 +2529,7 @@ async def test_multipage_image_ocrs_every_frame_end_to_end(self, isolated_env, m from lilbee.data.ingest import ingest_document - records = await ingest_document(f, "multi.tiff", "image") + records, _ = await ingest_document(f, "multi.tiff", "image") assert mock_svc.provider.vision_ocr.call_count == 3 assert sorted({r["page_start"] for r in records}) == [1, 2, 3] @@ -2675,7 +2699,7 @@ async def test_vision_backend_when_ocr_enabled_and_model_set( from lilbee.data.ingest import ingest_document - result = await ingest_document(f, "scanned.pdf", "pdf") + result, _ = await ingest_document(f, "scanned.pdf", "pdf") assert mock_svc.provider.pdf_ocr.call_args.kwargs["backend"] == "vision" assert mock_svc.provider.pdf_ocr.call_args.kwargs["per_page_timeout_s"] == 60.0 assert len(result) > 0 @@ -2698,7 +2722,7 @@ async def test_tesseract_backend_when_no_vision_model(self, mock_kf, isolated_en from lilbee.data.ingest import ingest_document - result = await ingest_document(f, "scanned.pdf", "pdf") + result, _ = await ingest_document(f, "scanned.pdf", "pdf") # Pool pdf_ocr is not used for tesseract. mock_svc.provider.pdf_ocr.assert_not_called() assert mock_kf.call_count == 2 @@ -2720,7 +2744,7 @@ async def test_tesseract_returning_empty_logs_skip( from lilbee.data.ingest import ingest_document with caplog.at_level("WARNING", logger="lilbee.data.ingest.extract"): - result = await ingest_document(f, "blank.pdf", "pdf") + result, _ = await ingest_document(f, "blank.pdf", "pdf") assert result == [] assert "Skipped blank.pdf" in caplog.text @@ -2756,7 +2780,7 @@ async def test_tesseract_no_timeout_runs_uncapped(self, mock_kf, isolated_env, m from lilbee.data.ingest import ingest_document - result = await ingest_document(f, "scanned.pdf", "pdf") + result, _ = await ingest_document(f, "scanned.pdf", "pdf") assert mock_kf.call_count == 2 assert len(result) > 0 @@ -2782,7 +2806,7 @@ async def _instant_timeout(coro, *, timeout): from lilbee.data.ingest import ingest_document with caplog.at_level("WARNING", logger="lilbee.data.ingest.extract"): - result = await ingest_document(f, "scanned.pdf", "pdf") + result, _ = await ingest_document(f, "scanned.pdf", "pdf") assert result == [] assert "Tesseract OCR exceeded" in caplog.text @@ -2803,7 +2827,7 @@ async def test_tesseract_extract_exception_logs_and_returns_empty( from lilbee.data.ingest import ingest_document with caplog.at_level("WARNING", logger="lilbee.data.ingest.extract"): - result = await ingest_document(f, "scanned.pdf", "pdf") + result, _ = await ingest_document(f, "scanned.pdf", "pdf") assert result == [] assert "OCR via tesseract backend failed" in caplog.text @@ -2868,7 +2892,7 @@ async def test_empty_markdown_returns_empty(self, isolated_env): md = isolated_env / "empty.md" md.write_text(" ") - result = await ingest_markdown(md, "empty.md") + result, _ = await ingest_markdown(md, "empty.md") assert result == [] async def test_no_chunks_returns_empty(self, isolated_env): @@ -2877,7 +2901,7 @@ async def test_no_chunks_returns_empty(self, isolated_env): md = isolated_env / "blank.md" md.write_text("some text") with mock.patch("lilbee.data.ingest.extract.chunk_text", return_value=[]): - result = await ingest_markdown(md, "blank.md") + result, _ = await ingest_markdown(md, "blank.md") assert result == [] async def test_frontmatter_only_produces_chunks(self, isolated_env): @@ -2885,7 +2909,7 @@ async def test_frontmatter_only_produces_chunks(self, isolated_env): md = isolated_env / "fm_only.md" md.write_text("---\ntitle: Just Frontmatter\ntags: [test]\n---\n") - result = await ingest_markdown(md, "fm_only.md") + result, _ = await ingest_markdown(md, "fm_only.md") assert len(result) > 0, "Frontmatter content should be indexed" @@ -2956,7 +2980,7 @@ async def test_empty_extraction_returns_empty(self, isolated_env): empty_result = mock.MagicMock(chunks=[]) mock_extract = Mock(return_value=empty_result) with mock.patch("kreuzberg.extract_file_sync", mock_extract): - result = await ingest_document(isolated_env / "e.xml", "e.xml", "xml") + result, _ = await ingest_document(isolated_env / "e.xml", "e.xml", "xml") assert result == [] async def test_no_chunks_returns_empty(self, isolated_env): @@ -2965,7 +2989,7 @@ async def test_no_chunks_returns_empty(self, isolated_env): no_chunks_result = mock.MagicMock(chunks=[]) mock_extract = Mock(return_value=no_chunks_result) with mock.patch("kreuzberg.extract_file_sync", mock_extract): - result = await ingest_document(isolated_env / "s.xml", "s.xml", "xml") + result, _ = await ingest_document(isolated_env / "s.xml", "s.xml", "xml") assert result == [] @@ -3233,12 +3257,21 @@ async def test_missing_metadata_falls_back_to_stem(self, mock_kf, isolated_env, assert item.meta.authors == "" assert all(r["title"] == "annual wildlife survey" for r in item.records) - async def test_markdown_title_derives_from_stem(self, isolated_env, mock_svc): - (isolated_env / "meeting_notes.md").write_text("# Heading\n\nSome content here.") + async def test_markdown_title_uses_the_h1_heading(self, isolated_env, mock_svc): + (isolated_env / "meeting_notes.md").write_text("# Project Kickoff\n\nSome content here.") + from lilbee.data.ingest import sync + + await sync(quiet=True) + items = mock_svc.store.write_chunks_batch.call_args.args[0] + item = next(it for it in items if it.source == "meeting_notes.md") + assert item.meta.title == "Project Kickoff" + assert all(r["title"] == "Project Kickoff" for r in item.records) + + async def test_markdown_without_h1_falls_back_to_stem(self, isolated_env, mock_svc): + (isolated_env / "meeting_notes.md").write_text("Some content, no heading.") from lilbee.data.ingest import sync await sync(quiet=True) items = mock_svc.store.write_chunks_batch.call_args.args[0] item = next(it for it in items if it.source == "meeting_notes.md") assert item.meta.title == "meeting notes" - assert all(r["title"] == "meeting notes" for r in item.records) diff --git a/tests/test_neighbors.py b/tests/test_neighbors.py index 63c103152..374916c4f 100644 --- a/tests/test_neighbors.py +++ b/tests/test_neighbors.py @@ -55,6 +55,12 @@ def test_no_overlap_joins_with_a_newline_seam(self): def test_fully_contained_text_adds_nothing(self): assert merge_adjacent_texts(["alpha beta", "beta"]) == "alpha beta" + def test_an_empty_side_has_no_overlap(self): + # An empty chunk shares nothing, so the seam is a plain join rather than + # a scan over a zero-length string. + assert merge_adjacent_texts(["", "delta"]) == "\ndelta" + assert merge_adjacent_texts(["alpha", ""]) == "alpha" + def test_merging_an_already_merged_passage_is_idempotent(self): texts = ["one two", "two three", "three four"] merged = merge_adjacent_texts(texts) diff --git a/tests/test_settings.py b/tests/test_settings.py index 615e05857..66eb017d5 100644 --- a/tests/test_settings.py +++ b/tests/test_settings.py @@ -4,6 +4,8 @@ import pytest +from lilbee.app.settings_map import SETTINGS_MAP, get_default +from lilbee.config_meta import WRITABLE_CONFIG_FIELDS from lilbee.core import settings @@ -207,14 +209,12 @@ def test_reranker_type_is_load_affecting(self): assert "reranker_type" in LOAD_AFFECTING_KEYS def test_reranker_fields_in_settings_map(self): - from lilbee.app.settings_map import SETTINGS_MAP assert "reranker_type" in SETTINGS_MAP 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 @@ -226,7 +226,6 @@ def test_fusion_knobs_in_settings_map(self): """The four adaptive-fusion / structural-filter knobs (which gate the on-by-default fusion behavior) are on the settings surface with their shipped defaults, so a dropped or typo'd entry fails CI.""" - from lilbee.app.settings_map import SETTINGS_MAP, get_default assert get_default("lexical_fusion_weight") == 1.0 assert get_default("adaptive_fusion") is True @@ -270,7 +269,6 @@ class TestMemoryTuningSettingsMap: """The dynamic-ctx tuning knobs are surfaced in the TUI settings map.""" def test_num_ctx_max_in_settings_map(self): - from lilbee.app.settings_map import SETTINGS_MAP, get_default defn = SETTINGS_MAP["num_ctx_max"] assert defn.writable is True @@ -279,7 +277,6 @@ def test_num_ctx_max_in_settings_map(self): assert get_default("num_ctx_max") is None def test_chat_n_ctx_target_in_settings_map(self): - from lilbee.app.settings_map import SETTINGS_MAP, get_default defn = SETTINGS_MAP["chat_n_ctx_target"] assert defn.writable is True @@ -292,7 +289,6 @@ def test_chat_n_ctx_target_in_settings_map(self): assert get_default("chat_n_ctx_target") == 8192 def test_flash_attention_in_settings_map(self): - from lilbee.app.settings_map import SETTINGS_MAP, get_default defn = SETTINGS_MAP["flash_attention"] assert defn.writable is True @@ -301,7 +297,6 @@ def test_flash_attention_in_settings_map(self): assert get_default("flash_attention") is None def test_kv_cache_type_in_settings_map(self): - from lilbee.app.settings_map import SETTINGS_MAP from lilbee.core.config.enums import KvCacheType defn = SETTINGS_MAP["kv_cache_type"] @@ -309,7 +304,6 @@ def test_kv_cache_type_in_settings_map(self): assert defn.choices == tuple(t.value for t in KvCacheType) def test_n_gpu_layers_in_settings_map(self): - from lilbee.app.settings_map import SETTINGS_MAP, get_default defn = SETTINGS_MAP["n_gpu_layers"] assert defn.writable is True @@ -317,7 +311,6 @@ def test_n_gpu_layers_in_settings_map(self): assert get_default("n_gpu_layers") is None def test_vision_ocr_max_tokens_in_settings_map(self): - from lilbee.app.settings_map import SETTINGS_MAP, get_default defn = SETTINGS_MAP["vision_ocr_max_tokens"] assert defn.writable is True @@ -327,7 +320,6 @@ def test_vision_ocr_max_tokens_in_settings_map(self): assert get_default("vision_ocr_max_tokens") == 4096 def test_vision_ocr_concurrency_in_settings_map(self): - from lilbee.app.settings_map import SETTINGS_MAP, get_default defn = SETTINGS_MAP["vision_ocr_concurrency"] assert defn.writable is True @@ -337,7 +329,6 @@ def test_vision_ocr_concurrency_in_settings_map(self): assert get_default("vision_ocr_concurrency") == 4 def test_crawl_render_mode_in_settings_map(self): - from lilbee.app.settings_map import SETTINGS_MAP from lilbee.core.config.enums import CrawlRenderMode defn = SETTINGS_MAP["crawl_render_mode"] @@ -346,14 +337,12 @@ def test_crawl_render_mode_in_settings_map(self): assert defn.choices == tuple(m.value for m in CrawlRenderMode) def test_crawl_render_mode_is_writable_for_programmatic_surfaces(self): - from lilbee.config_meta import WRITABLE_CONFIG_FIELDS # The TUI checkbox persists the choice via apply_settings_update, so the # field must be writable through the HTTP / MCP / programmatic contract. assert "crawl_render_mode" in WRITABLE_CONFIG_FIELDS def test_browser_memory_levers_in_settings_map(self): - from lilbee.app.settings_map import SETTINGS_MAP, get_default recycle = SETTINGS_MAP["crawl_browser_recycle_pages"] assert recycle.writable is True @@ -485,19 +474,16 @@ def test_auto_sync_defaults_true(self): assert Config().auto_sync is True def test_auto_sync_is_writable(self): - from lilbee.config_meta import WRITABLE_CONFIG_FIELDS assert "auto_sync" in WRITABLE_CONFIG_FIELDS def test_auto_sync_in_settings_map(self): - from lilbee.app.settings_map import SETTINGS_MAP assert "auto_sync" in SETTINGS_MAP class TestListSettingRegexMarker: def test_only_regex_list_validates_as_regex(self): - from lilbee.app.settings_map import SETTINGS_MAP assert SETTINGS_MAP["crawl_exclude_patterns"].validate_regex is True # Chromium flag list must not be regex-validated. @@ -530,7 +516,6 @@ class TestTitleSearchSettings: """The title-arm knobs are exposed on every settings surface.""" def test_title_search_in_settings_map(self): - from lilbee.app.settings_map import SETTINGS_MAP, get_default defn = SETTINGS_MAP["title_search"] assert defn.writable is True @@ -539,7 +524,6 @@ def test_title_search_in_settings_map(self): assert get_default("title_search") is False def test_title_search_weight_in_settings_map(self): - from lilbee.app.settings_map import SETTINGS_MAP, get_default defn = SETTINGS_MAP["title_search_weight"] assert defn.writable is True @@ -548,7 +532,6 @@ def test_title_search_weight_in_settings_map(self): assert get_default("title_search_weight") == 0.5 def test_title_search_fields_are_writable_for_programmatic_surfaces(self): - from lilbee.config_meta import WRITABLE_CONFIG_FIELDS assert "title_search" in WRITABLE_CONFIG_FIELDS assert "title_search_weight" in WRITABLE_CONFIG_FIELDS diff --git a/tests/test_store.py b/tests/test_store.py index 6cfe88941..5306516de 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -11,6 +11,7 @@ CitationRecord, SearchChunk, SearchScope, + SourceMeta, SourceType, Store, cosine_sim, @@ -2774,8 +2775,6 @@ class TestSourceMetadata: """Extraction-time document metadata persisted on the sources table.""" def test_upsert_source_persists_meta(self, store): - from lilbee.data.store import SourceMeta - store.upsert_source( "a.pdf", "hash1", @@ -2798,8 +2797,6 @@ def test_pre_meta_sources_table_evolves_in_place(self, store): """An old sources table gains the metadata columns on the next write.""" import pyarrow as pa - from lilbee.data.store import SourceMeta - old_schema = pa.schema( [ pa.field("filename", pa.utf8()), From eb39ab6a9979e433177f41d8145516e87ddd7876 Mon Sep 17 00:00:00 2001 From: Tobias Perelstein <5562156+tobocop2@users.noreply.github.com> Date: Mon, 20 Jul 2026 10:26:08 -0400 Subject: [PATCH 23/28] Do not let an unresolvable home take down config loading The data-root expanduser() added for the "~/lilbee" env case was unguarded, but expanduser() raises when the OS cannot resolve a home directory (no HOME, an unknown ~user). Configuration must still load in that case, so an unexpandable path now keeps its literal form. This also fixes CI: a test that monkeypatches Path.expanduser to raise made every Config construction during its teardown blow up, erroring the whole test job on all platforms while main stayed green. --- src/lilbee/core/config/model.py | 17 ++++++++++++++++- tests/test_config.py | 12 ++++++++++++ 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/src/lilbee/core/config/model.py b/src/lilbee/core/config/model.py index 104cbb4f1..3cf807fc2 100644 --- a/src/lilbee/core/config/model.py +++ b/src/lilbee/core/config/model.py @@ -44,6 +44,19 @@ _UNSET_PATH = Path() +def _expanded(path: Path) -> Path: + """``path`` with a leading ``~`` expanded, or unchanged if it cannot be. + + ``expanduser()`` raises when the OS cannot resolve a home directory (no + HOME, an unknown ``~user``). Configuration must still load in that case, so + the literal path stands rather than taking the whole process down. + """ + try: + return path.expanduser() + except (OSError, RuntimeError): + return path + + class Config(BaseSettings): """Runtime configuration: one singleton instance, mutated by CLI overrides.""" @@ -1040,7 +1053,9 @@ def _resolve_defaults(cls, data: Any) -> Any: # expanduser() so a "~/lilbee" value from a systemd unit or .env (which # do not expand ~) points at the home dir instead of creating a literal # ./~ tree that a server lock keyed on this path would then diverge on. - root = Path(data["data_root"]).expanduser() + # It raises when the OS cannot resolve a home directory, and config must + # still load then, so an unresolvable ~ keeps the literal path. + root = _expanded(Path(data["data_root"])) data["data_root"] = root if data.get("documents_dir") in (None, _UNSET_PATH): data["documents_dir"] = root / "documents" diff --git a/tests/test_config.py b/tests/test_config.py index 59257d58a..3264323f0 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -135,6 +135,18 @@ def test_empty_data_root_falls_back_to_default_not_cwd(self): assert c2.data_root != Path.cwd() assert str(c2.data_root).endswith("lilbee") + def test_unexpandable_home_keeps_the_literal_path(self): + """expanduser() raises when the OS cannot resolve a home directory (no + HOME, an unknown ~user). The data root must still resolve rather than + taking config construction -- and the whole process -- down.""" + from lilbee.core.config.model import _expanded + + literal = Path("~/lilbee") + with mock.patch.object(Path, "expanduser", side_effect=OSError("no home")): + assert _expanded(literal) == literal + # The normal case still expands. + assert _expanded(literal) == Path.home() / "lilbee" + def test_local_server_urls_from_env(self, tmp_path): env = _clean_env(tmp_path) env["LILBEE_OLLAMA_BASE_URL"] = "http://box:11434" From af56a5d90d5026d84a5d64faa6726a43c21dfd71 Mon Sep 17 00:00:00 2001 From: Tobias Perelstein <5562156+tobocop2@users.noreply.github.com> Date: Mon, 20 Jul 2026 11:07:34 -0400 Subject: [PATCH 24/28] Drop the dead budget-fit wrapper and fix overlap on control bytes The context-budget fit kept a wrapper that no production caller reached after neighbor expansion started deriving its leftover from the fit; only tests called it, so seven of them asserted against a shim rather than the path that runs. Deleted it and pointed the tests at the real pair. The neighbor overlap scan joined the two chunks around a NUL sentinel, which is only sound for text that carries none of its own. Extracted document text does carry stray control bytes, and there the match ran past the right chunk's end and merging dropped the tail it claimed to share. It now runs the prefix function over the right chunk alone and streams the left chunk's tail through it, so no sentinel is involved. Checked against a brute-force reference over 300k random cases. Also: an image skipped because OCR is off no longer pays for a metadata extraction it cannot use, and the markdown H1 scan drops a heading test that could never be false. --- src/lilbee/data/ingest/extract.py | 11 +++++++---- src/lilbee/retrieval/query/neighbors.py | 26 +++++++++++++++++-------- src/lilbee/retrieval/query/searcher.py | 26 ++++++------------------- tests/test_ingest.py | 11 +++++++++-- tests/test_neighbors.py | 6 ++++++ tests/test_query.py | 25 +++++++++++++++++------- 6 files changed, 64 insertions(+), 41 deletions(-) diff --git a/src/lilbee/data/ingest/extract.py b/src/lilbee/data/ingest/extract.py index ca59f935a..4474364e4 100644 --- a/src/lilbee/data/ingest/extract.py +++ b/src/lilbee/data/ingest/extract.py @@ -539,12 +539,15 @@ async def _handle_image( Its EXIF/XMP metadata is captured separately so the title/authors/date are not lost to the stem fallback. """ - meta = await to_ingest_thread(_image_meta, path, source_name) if _effective_enable_ocr() is False: # OCR explicitly disabled: an image has no text layer, so skip it rather - # than paying the full Tesseract cost the config says is turned off. + # than paying the full Tesseract cost the config says is turned off. The + # metadata read is skipped with it -- a file that contributes no text + # needs no title beyond its stem, and an image-heavy library would pay + # one extraction per skipped file for nothing. log.info("OCR disabled; skipping image OCR for %s", source_name) - return [], meta + return [], SourceMeta(title=derive_title(source_name)) + meta = await to_ingest_thread(_image_meta, path, source_name) vision_model = active_config().vision_model if _should_run_ocr() and vision_model: log.info("Image: using vision OCR for %s (model=%s)", source_name, vision_model) @@ -654,7 +657,7 @@ def _markdown_h1(text: str) -> str | None: stripped = line.strip() if not stripped: continue - if stripped.startswith("# ") and not stripped.startswith("## "): + if stripped.startswith("# "): return stripped[2:].strip() or None return None return None diff --git a/src/lilbee/retrieval/query/neighbors.py b/src/lilbee/retrieval/query/neighbors.py index 11eade2be..bc4a64ec5 100644 --- a/src/lilbee/retrieval/query/neighbors.py +++ b/src/lilbee/retrieval/query/neighbors.py @@ -32,18 +32,28 @@ def _overlap_chars(left: str, right: str) -> int: """ if not left or not right: return 0 - # Only left's last len(right) chars can match a prefix of right, so the - # probe stays linear in the incoming chunk, not the accumulated passage. - probe = f"{right}\0{left[-len(right) :]}" - failure = [0] * len(probe) + # Prefix function of right alone. Concatenating the two around a sentinel + # would be shorter to write but only correct for text that never contains + # the sentinel, and extracted document text carries stray control bytes. + failure = [0] * len(right) length = 0 - for i in range(1, len(probe)): - while length and probe[i] != probe[length]: + for i in range(1, len(right)): + while length and right[i] != right[length]: length = failure[length - 1] - if probe[i] == probe[length]: + if right[i] == right[length]: length += 1 failure[i] = length - return failure[-1] + # Only left's last len(right) chars can reach a prefix of right, so the scan + # stays linear in the incoming chunk, not the accumulated passage. That + # window is also what keeps ``right[length]`` in range: matching right in + # full costs len(right) chars, so it can only complete on the last one. + length = 0 + for char in left[-len(right) :]: + while length and char != right[length]: + length = failure[length - 1] + if char == right[length]: + length += 1 + return length def merge_adjacent_texts(texts: list[str]) -> str: diff --git a/src/lilbee/retrieval/query/searcher.py b/src/lilbee/retrieval/query/searcher.py index f004b0329..7ae18bbf5 100644 --- a/src/lilbee/retrieval/query/searcher.py +++ b/src/lilbee/retrieval/query/searcher.py @@ -983,29 +983,15 @@ def _context_budget( ) return int((prompt_token_budget(ctx) - non_source) * 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. - """ - return self._fit_to_budget(results, self._context_budget(system, question, history, scale))[ - 0 - ] - def _fit_to_budget( self, results: list[SearchChunk], budget: int ) -> tuple[list[SearchChunk], int]: """Fit *results* into *budget*: the kept sources and the tokens they cost. + ``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. + Returning the spent total lets the caller derive the leftover for neighbor expansion instead of re-deriving the same per-chunk cost, so the two stages cannot drift apart on the accounting. @@ -1415,8 +1401,8 @@ def ask_stream( return results, messages = rag.results, rag.messages # No overflow retry here: a stream cannot be rebuilt once tokens have - # been yielded, so the conservative budget in _fit_context_budget is - # the streaming path's protection. + # been yielded, so the conservative budget the context fit already + # applied is the streaming path's protection. provider_messages = self._messages_for_provider(messages) opts = options if options is not None else self._config.generation_options() events = stream_chat_with_cap( diff --git a/tests/test_ingest.py b/tests/test_ingest.py index 3d28da7fb..816b8a29b 100644 --- a/tests/test_ingest.py +++ b/tests/test_ingest.py @@ -2470,11 +2470,18 @@ async def test_image_ocr_disabled_skips_both_backends(self, isolated_env, mock_s from lilbee.data.ingest import ingest_document - with mock.patch("lilbee.data.ingest.extract._run_tesseract_sync") as mock_tess: - result, _ = await ingest_document(f, "scan.png", "image") + with ( + mock.patch("lilbee.data.ingest.extract._run_tesseract_sync") as mock_tess, + mock.patch("lilbee.data.ingest.extract._image_meta") as mock_meta, + ): + result, meta = await ingest_document(f, "scan.png", "image") assert result == [] mock_svc.provider.vision_ocr.assert_not_called() mock_tess.assert_not_called() + # A skipped image contributes no text, so it must not pay for a metadata + # extraction either: an image-heavy library would run one per file. + mock_meta.assert_not_called() + assert meta.title == "scan" async def test_vision_ocr_cache_key_includes_timeout(self, isolated_env, mock_svc, monkeypatch): """The vision OCR cache key carries the per-page timeout so raising it diff --git a/tests/test_neighbors.py b/tests/test_neighbors.py index 374916c4f..1613b4a9d 100644 --- a/tests/test_neighbors.py +++ b/tests/test_neighbors.py @@ -61,6 +61,12 @@ def test_an_empty_side_has_no_overlap(self): assert merge_adjacent_texts(["", "delta"]) == "\ndelta" assert merge_adjacent_texts(["alpha", ""]) == "alpha" + def test_control_bytes_in_the_text_do_not_confuse_the_overlap(self): + # Extracted document text can carry stray control bytes. The shared + # region here is the trailing "\0a\0"; a scan that joined the two sides + # around a NUL sentinel would match past it and swallow right's tail. + assert merge_adjacent_texts(["Xa\0a\0", "\0a\0a"]) == "Xa\0a\0a" + def test_merging_an_already_merged_passage_is_idempotent(self): texts = ["one two", "two three", "three four"] merged = merge_adjacent_texts(texts) diff --git a/tests/test_query.py b/tests/test_query.py index c8155041d..cbb38a11c 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -136,6 +136,17 @@ def _make_result( ) +def _fit( + results: list[SearchChunk], + system: str = "sys", + question: str = "q", + history: list | None = None, +) -> list[SearchChunk]: + """Sources kept by the budget fit, derived exactly as _finalize_context does.""" + searcher = get_services().searcher + return searcher._fit_to_budget(results, searcher._context_budget(system, question, history))[0] + + class TestDisplaySourcePath: """source citations render absolute paths with ~ expansion.""" @@ -882,7 +893,7 @@ def _restore_ctx(self): def test_trims_lowest_ranked_sources_to_fit(self, mock_svc): cfg.num_ctx = 1400 results = [_make_result(source=f"{i}.pdf", chunk="x" * 300) for i in range(5)] - kept = get_services().searcher._fit_context_budget(results, "sys", "q", None) + kept = _fit(results) assert 0 < len(kept) < len(results) assert kept == results[: len(kept)] # keeps the top-ranked prefix @@ -900,7 +911,7 @@ def test_fitted_context_survives_the_provider_that_enforces_the_window(self, moc system, question = "sys " * 40, "q " * 20 results = [_make_result(source=f"{i}.pdf", chunk="x" * 900) for i in range(5)] - kept = get_services().searcher._fit_context_budget(results, system, question, None) + kept = _fit(results, system, question) searcher = get_services().searcher assembled = ( @@ -917,12 +928,12 @@ def test_fitted_context_survives_the_provider_that_enforces_the_window(self, moc def test_keeps_all_when_budget_ample(self, mock_svc): cfg.num_ctx = 100_000 results = [_make_result(source=f"{i}.pdf", chunk="short") for i in range(5)] - assert get_services().searcher._fit_context_budget(results, "sys", "q", None) == results + assert _fit(results) == results def test_keeps_top_source_even_if_alone_over_budget(self, mock_svc): cfg.num_ctx = 1 results = [_make_result(source="big.pdf", chunk="x" * 9000), _make_result(source="b.pdf")] - kept = get_services().searcher._fit_context_budget(results, "sys", "q", None) + kept = _fit(results) assert len(kept) == 1 def test_history_counts_against_the_budget(self, mock_svc): @@ -932,15 +943,15 @@ def test_history_counts_against_the_budget(self, mock_svc): cfg.num_ctx = 3000 results = [_make_result(source=f"{i}.pdf", chunk="x" * 1200) for i in range(5)] history = [{"role": "user", "content": "h" * 3000}] - no_hist = get_services().searcher._fit_context_budget(results, "sys", "q", None) - with_hist = get_services().searcher._fit_context_budget(results, "sys", "q", history) + no_hist = _fit(results) + with_hist = _fit(results, history=history) assert len(with_hist) < len(no_hist) def test_logs_when_trimming(self, mock_svc, caplog): cfg.num_ctx = 1400 results = [_make_result(source=f"{i}.pdf", chunk="x" * 300) for i in range(5)] with caplog.at_level("INFO"): - get_services().searcher._fit_context_budget(results, "sys", "q", None) + _fit(results) assert "to fit the model context window" in caplog.text From b245948b15c0fefe115d5a731dafdf80a2e453df Mon Sep 17 00:00:00 2001 From: Tobias Perelstein <5562156+tobocop2@users.noreply.github.com> Date: Mon, 20 Jul 2026 11:15:52 -0400 Subject: [PATCH 25/28] Go back to the simple overlap scan; the clever one was slower The neighbor overlap scan was rewritten during this branch into a linear prefix-function pass, on the reasoning that trying every length would be quadratic when adjacent chunks share nothing. Measured, that premise is wrong: endswith rejects on its first differing character in C, so the simple scan is within ~1.3x on random text and hundreds of times faster on repetitive text, where the longest length matches immediately. Both are far under a millisecond on a path that runs a handful of times per query. So the simple version wins on every count: fewer lines, correct by inspection, and it cannot have the sentinel-collision bug the linear version shipped. Verified identical to the previous behaviour over 200k random cases including the control bytes that broke the sentinel. --- src/lilbee/retrieval/query/neighbors.py | 35 ++++++------------------- 1 file changed, 8 insertions(+), 27 deletions(-) diff --git a/src/lilbee/retrieval/query/neighbors.py b/src/lilbee/retrieval/query/neighbors.py index bc4a64ec5..d456a083c 100644 --- a/src/lilbee/retrieval/query/neighbors.py +++ b/src/lilbee/retrieval/query/neighbors.py @@ -26,34 +26,15 @@ def _overlap_chars(left: str, right: str) -> int: fully contained text matches whole, which is what makes merging an already-widened passage idempotent instead of duplicating its neighbors. - Computed with a prefix-function pass rather than trying every length: the - no-overlap case (a ``chunk_overlap=0`` build) is common and would otherwise - scan every length to exhaustion, quadratic in the chunk's own size. + Longest length first, so the first match wins. A prefix-function scan is + the better complexity on paper but measured slower here on every input + shape tried: each ``endswith`` rejects on its first differing character in + C, while the linear version pays per-character interpreter overhead. """ - if not left or not right: - return 0 - # Prefix function of right alone. Concatenating the two around a sentinel - # would be shorter to write but only correct for text that never contains - # the sentinel, and extracted document text carries stray control bytes. - failure = [0] * len(right) - length = 0 - for i in range(1, len(right)): - while length and right[i] != right[length]: - length = failure[length - 1] - if right[i] == right[length]: - length += 1 - failure[i] = length - # Only left's last len(right) chars can reach a prefix of right, so the scan - # stays linear in the incoming chunk, not the accumulated passage. That - # window is also what keeps ``right[length]`` in range: matching right in - # full costs len(right) chars, so it can only complete on the last one. - length = 0 - for char in left[-len(right) :]: - while length and char != right[length]: - length = failure[length - 1] - if char == right[length]: - length += 1 - return length + 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: From d9bd8678dbb6b75be2f3221f9d47f6455f868253 Mon Sep 17 00:00:00 2001 From: Tobias <5562156+tobocop2@users.noreply.github.com> Date: Mon, 20 Jul 2026 11:45:57 -0400 Subject: [PATCH 26/28] Stop the chunk fetch from loading the whole index into memory (#567) * Drop the whole-table Arrow fallback from the two chunk-fetch paths get_chunks_by_source and get_chunks_by_indices rescued a failed filtered query by calling table.to_arrow(), which materializes every chunk of every document -- embedding vectors included -- before filtering in Python. On a large library that is a multi-gigabyte spike, and neighbor expansion hits it once per hit source per query. The fallback guarded a LanceDB version whose FTS-indexed tables rejected .where() on arbitrary columns. That no longer reproduces: a filtered empty-search query returns the right rows on an FTS-indexed table, so the path was dead code held alive by tests that mocked search() into raising. Both functions now let a query failure propagate. An exception names the problem; an OOM kill from the rescue scan leaves nothing to debug. The two mocked-fallback tests are replaced by tests that build the FTS index for real and assert the filtered query still selects rows, so a future LanceDB bump that breaks the predicate fails CI instead of failing in production. * Stop the autocomplete OSError test from leaking into config validation The test patched pathlib.Path.expanduser globally; the config validator also expands the data root on every field assignment, so the cfg-restoring fixture raised through the patch at teardown and turned the whole suite red. Patch only the autocomplete module's Path instead. The breakage predates this branch (it came with the data-root expanduser fix on the base) but blocks this PR's CI, so carry the one-test fix here too. --- src/lilbee/data/store/core.py | 41 +++++++++--------------- tests/test_store.py | 59 ++++++++++------------------------- tests/test_tui_widgets.py | 20 ++++++++---- 3 files changed, 44 insertions(+), 76 deletions(-) diff --git a/src/lilbee/data/store/core.py b/src/lilbee/data/store/core.py index 816e65226..e210a933e 100644 --- a/src/lilbee/data/store/core.py +++ b/src/lilbee/data/store/core.py @@ -11,7 +11,6 @@ from typing import TYPE_CHECKING import pyarrow as pa -import pyarrow.compute as pc from lilbee.core.config import ( CHUNK_CONCEPTS_TABLE, @@ -1098,29 +1097,28 @@ def count_chunks(self) -> int: return table.count_rows() if table is not None else 0 def get_chunks_by_source(self, source: str) -> list[SearchChunk]: - """Return every chunk whose ``source`` equals *source*.""" + """Return every chunk whose ``source`` equals *source*. + + The database does the filtering, so only the matching rows are read. + A query failure raises rather than falling back to a whole-table scan: + a document's chunks are a bounded read, and the scan that would rescue + it costs the entire index, vectors included, in memory. + """ table = self.open_table(CHUNKS_TABLE) if table is None: return [] escaped = escape_sql_string(source) - try: - rows = table.search().where(f"source = '{escaped}'").limit(None).to_list() - except Exception: - # FTS-enabled tables return a query builder that cannot - # handle .where() on arbitrary columns; fall through to a - # pyarrow.compute filter on the Arrow table so the source - # match runs in C++ without materializing non-matching rows. - log.debug("get_chunks_by_source search() failed, using Arrow fallback", exc_info=True) - arrow_tbl = table.to_arrow() - filtered = arrow_tbl.filter(pc.equal(arrow_tbl["source"], source)) - rows = filtered.to_pylist() + rows = table.search().where(f"source = '{escaped}'").limit(None).to_list() 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. + the document are simply absent from the result. Filtering happens in + the database for the same reason as :meth:`get_chunks_by_source`: + neighbor expansion runs once per hit source per query, so a + whole-table rescue would spike memory on the hottest path there is. """ if not indices: return [] @@ -1129,19 +1127,8 @@ def get_chunks_by_indices(self, source: str, indices: Sequence[int]) -> list[Sea 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() + predicate = f"source = '{escaped}' AND chunk_index IN ({wanted})" + rows = table.search().where(predicate).limit(None).to_list() return sorted((SearchChunk(**r) for r in rows), key=lambda c: c.chunk_index) def _delete_by_sources_unlocked(self, sources: list[str]) -> None: diff --git a/tests/test_store.py b/tests/test_store.py index 5306516de..642938629 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -1482,31 +1482,17 @@ def test_chunk_type_defaults_to_raw(self, store): results = store.get_chunks_by_source("doc0.md") assert results[0].chunk_type == "raw" - def test_get_chunks_by_source_fallback(self, store): - """Fallback path when table.search() raises (e.g. incompatible FTS builder).""" - from unittest.mock import patch + def test_get_chunks_by_source_filters_with_fts_index_built(self, store): + """The filtered query still selects rows once the chunks table is FTS-indexed. - records = _make_records(n=2) - store.add_chunks(records) - - # Make table.search() raise to trigger the Arrow fallback - original_open = store.open_table - - def _broken_open(name): - table = original_open(name) - if table is None: - return None - - 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_source("doc0.md") - assert len(results) == 1 - assert results[0].source == "doc0.md" + Both chunk-fetch paths rely on the database doing the filtering, so an + FTS-indexed table that rejected ``.where()`` would silently regress them + into whole-table reads. Pin the behavior the fetch paths depend on. + """ + store.add_chunks(_make_records(n=3)) + store.ensure_fts_index() + results = store.get_chunks_by_source("doc1.md") + assert [r.source for r in results] == ["doc1.md"] def _one_source_records(source: str, n: int) -> list[dict]: @@ -1541,26 +1527,13 @@ def test_empty_indices_returns_empty(self, store): 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]) + def test_filters_with_fts_index_built(self, store): + """The compound source+index predicate survives an FTS-indexed table.""" + store.add_chunks(_one_source_records("a.md", 3) + _one_source_records("b.md", 3)) + store.ensure_fts_index() + results = store.get_chunks_by_indices("a.md", [0, 2]) assert [r.chunk_index for r in results] == [0, 2] + assert {r.source for r in results} == {"a.md"} def test_search_chunk_default_is_raw(self): chunk = SearchChunk( diff --git a/tests/test_tui_widgets.py b/tests/test_tui_widgets.py index 0f141bc7e..c99a40419 100644 --- a/tests/test_tui_widgets.py +++ b/tests/test_tui_widgets.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import Any +from typing import Any, Never from unittest import mock import pytest @@ -2755,15 +2755,23 @@ def test_add_arg_completions(self, tmp_path: object) -> None: assert any("testfile.txt" in x for x in r) def test_path_exists_swallows_oserror(self, monkeypatch: pytest.MonkeyPatch) -> None: - """A path the OS refuses to stat must read as missing, not raise.""" - from pathlib import Path as P + """A path the OS refuses to stat must read as missing, not raise. + Patches this module's ``Path`` rather than ``pathlib.Path.expanduser`` + itself: a global patch also breaks the config validator, which expands + the data root on every field assignment, and so blows up in the + cfg-restoring fixture's teardown before monkeypatch unwinds. + """ from lilbee.cli.tui.widgets import autocomplete - def _boom(self: P) -> P: - raise OSError("bad path") + class _BoomPath: + def __init__(self, *_args: object) -> None: + pass + + def expanduser(self) -> Never: + raise OSError("bad path") - monkeypatch.setattr(P, "expanduser", _boom) + monkeypatch.setattr(autocomplete, "Path", _BoomPath) assert autocomplete._path_exists("~oops") is False def test_add_complete_path_collapses_so_enter_submits(self, tmp_path: object) -> None: From 2dffa93e996bd2522200afee74085557a139771c Mon Sep 17 00:00:00 2001 From: Tobias Perelstein <5562156+tobocop2@users.noreply.github.com> Date: Mon, 20 Jul 2026 12:11:09 -0400 Subject: [PATCH 27/28] Tidy three retrieval seams: spaCy import, HyDE budget, dead compaction slot The ingest pipeline and the entity backfill both reached across packages for concepts.nlp's private _ensure_spacy_model. Both now call load_spacy_pipeline, the public wrapper that exists for exactly this. concepts/graph.py keeps the private call: it lives in the same package. HyDE was capped by EXPANSION_MAX_TOKENS, which also caps query-variant expansion. A hypothetical answer passage and a list of query rewrites are different lengths of output, so retuning one was silently retuning the other. HyDE gets its own constant. COMPACT_PROMPT carried a {previous} slot for carried-forward notes, and batch_overflow reserved window space for them, but every caller passes an empty string: summarize_history merges per-batch notes instead of feeding them back in (deliberately, to keep summary depth at one). The prompt scaffolding never rendered and the reservation only made batches smaller than they needed to be. Both are gone, along with the previous_summary parameter that threaded through plan_compaction, batch_overflow, and _summarize_batch. --- src/lilbee/data/ingest/pipeline.py | 4 +-- src/lilbee/retrieval/entities/lifecycle.py | 4 +-- src/lilbee/retrieval/query/compaction.py | 20 +++--------- src/lilbee/retrieval/query/expansion.py | 7 ++++- src/lilbee/retrieval/query/searcher.py | 36 ++++++++++++---------- tests/test_compaction.py | 16 +++------- tests/test_entities_ingest.py | 4 +-- tests/test_entities_lifecycle.py | 4 +-- tests/test_query.py | 11 +++++++ tests/test_summarize_history.py | 2 +- 10 files changed, 55 insertions(+), 53 deletions(-) diff --git a/src/lilbee/data/ingest/pipeline.py b/src/lilbee/data/ingest/pipeline.py index b6035580f..b21419b85 100644 --- a/src/lilbee/data/ingest/pipeline.py +++ b/src/lilbee/data/ingest/pipeline.py @@ -143,11 +143,11 @@ async def _build_entity_records(records: list[ChunkRecord], source_name: str) -> nlp = None if any(t.kind is ExtractorKind.SPACY for t in schema.types): from lilbee.retrieval.concepts import concepts_available - from lilbee.retrieval.concepts.nlp import _ensure_spacy_model + from lilbee.retrieval.concepts.nlp import load_spacy_pipeline if concepts_available(): try: - nlp = _ensure_spacy_model() + nlp = load_spacy_pipeline() except ImportError: log.warning("spaCy model unavailable; spacy-kind entity types skipped") provider = None diff --git a/src/lilbee/retrieval/entities/lifecycle.py b/src/lilbee/retrieval/entities/lifecycle.py index b1e72e93f..afeaa13c6 100644 --- a/src/lilbee/retrieval/entities/lifecycle.py +++ b/src/lilbee/retrieval/entities/lifecycle.py @@ -197,10 +197,10 @@ def _full_pass(store: Store, schema: EntitySchema, cancel: threading.Event | Non from lilbee.retrieval.concepts import concepts_available if concepts_available(): - from lilbee.retrieval.concepts.nlp import _ensure_spacy_model + from lilbee.retrieval.concepts.nlp import load_spacy_pipeline try: - nlp = _ensure_spacy_model() + nlp = load_spacy_pipeline() except ImportError: log.warning("spaCy model unavailable; skipping spaCy entity types") provider = None diff --git a/src/lilbee/retrieval/query/compaction.py b/src/lilbee/retrieval/query/compaction.py index 79c961acd..59ab01645 100644 --- a/src/lilbee/retrieval/query/compaction.py +++ b/src/lilbee/retrieval/query/compaction.py @@ -14,7 +14,6 @@ from lilbee.retrieval.query.history_window import ( chars_for_tokens, - estimate_text_tokens, estimate_tokens, windowed_history, ) @@ -28,7 +27,7 @@ "Condense the conversation below into brief factual notes that let an " "assistant carry it on. Keep names, numbers, decisions, and anything left " "unresolved; drop pleasantries. Under {words} words. Return ONLY the notes.\n\n" - "{previous}Conversation:\n{transcript}" + "Conversation:\n{transcript}" ) # Ceiling on a summary's tokens; ctx/8 governs below it. A tighter ceiling @@ -204,9 +203,7 @@ def summary_cap(ctx_target: int) -> int: return max(_SUMMARY_MIN_TOKENS, min(COMPACT_MAX_TOKENS, ctx_target // _SUMMARY_CTX_FRACTION)) -def batch_overflow( - dropped: list[ChatMessage], previous_summary: str, *, ctx_target: int -) -> list[list[ChatMessage]]: +def batch_overflow(dropped: list[ChatMessage], *, ctx_target: int) -> list[list[ChatMessage]]: """Split *dropped* into batches whose summarize prompt each fit *ctx_target*. Compaction usually nibbles a pair at a time, but switching models does not: @@ -226,12 +223,7 @@ def batch_overflow( room = max( _MIN_BATCH_TOKENS, int( - ( - ctx_target - - summary_cap(ctx_target) - - estimate_text_tokens(previous_summary) - - _PROMPT_OVERHEAD_TOKENS - ) + (ctx_target - summary_cap(ctx_target) - _PROMPT_OVERHEAD_TOKENS) * _ESTIMATE_SAFETY_FRACTION ), ) @@ -256,9 +248,7 @@ def batch_overflow( return batches -def plan_compaction( - dropped: list[ChatMessage], previous_summary: str, *, ctx_target: int -) -> CompactionPlan: +def plan_compaction(dropped: list[ChatMessage], *, ctx_target: int) -> CompactionPlan: """Decide what one compaction condenses, keeping its cost bounded. Beyond MAX_COMPACT_CALLS batches the oldest turns are stranded rather than @@ -267,7 +257,7 @@ def plan_compaction( trade. The most recent slice is the part still worth remembering, and the caller reports the stranded count instead of pretending it survived. """ - batches = batch_overflow(dropped, previous_summary, ctx_target=ctx_target) + batches = batch_overflow(dropped, ctx_target=ctx_target) if len(batches) <= MAX_COMPACT_CALLS: return CompactionPlan(batches=batches, stranded=0) return CompactionPlan( diff --git a/src/lilbee/retrieval/query/expansion.py b/src/lilbee/retrieval/query/expansion.py index ba03126b3..d8f83b546 100644 --- a/src/lilbee/retrieval/query/expansion.py +++ b/src/lilbee/retrieval/query/expansion.py @@ -1,4 +1,4 @@ -"""Constants for LLM-driven query expansion and history condensation.""" +"""Constants for LLM-driven query expansion, HyDE, and history condensation.""" from __future__ import annotations @@ -10,6 +10,11 @@ EXPANSION_MAX_TOKENS = 200 +# HyDE writes one hypothetical answer passage to embed, which is a different +# shape of output from a list of query variants. Same number today, but tuning +# either one must not move the other. +HYDE_MAX_TOKENS = 200 + CONDENSE_PROMPT = ( "Rewrite the follow-up question as one standalone search query, resolving " "pronouns and references using the conversation. Return ONLY the rewritten " diff --git a/src/lilbee/retrieval/query/searcher.py b/src/lilbee/retrieval/query/searcher.py index 7ae18bbf5..2dd4df416 100644 --- a/src/lilbee/retrieval/query/searcher.py +++ b/src/lilbee/retrieval/query/searcher.py @@ -52,6 +52,7 @@ CONDENSE_PROMPT, EXPANSION_MAX_TOKENS, EXPANSION_PROMPT, + HYDE_MAX_TOKENS, ) from lilbee.retrieval.query.formatting import ( CONTEXT_TEMPLATE, @@ -431,7 +432,7 @@ def _hyde_search( response = self._provider.chat( [{"role": "user", "content": self._config.hyde_prompt.format(question=question)}], stream=False, - options={"num_predict": EXPANSION_MAX_TOKENS}, + options={"num_predict": HYDE_MAX_TOKENS}, ) # Reasoning models front-load deliberation; embedding it instead # of the passage would search for the model's thought process. @@ -700,14 +701,14 @@ def summarize_history( ``on_batch`` hears ``(batch, total)`` before each model call, for progress UI. """ ctx_target = self._config.chat_n_ctx_target - plan = plan_compaction(messages, previous_summary, ctx_target=ctx_target) + plan = plan_compaction(messages, ctx_target=ctx_target) notes: list[str] = [] condensed = 0 stranded = plan.stranded for index, batch in enumerate(plan.batches): if on_batch is not None: on_batch(index + 1, len(plan.batches)) - note = self._summarize_batch(batch, "") + note = self._summarize_batch(batch) if note: notes.append(note) condensed += len(batch) @@ -718,26 +719,29 @@ def summarize_history( merged = merge_notes(previous_summary, notes) cap = summary_cap(ctx_target) if estimate_text_tokens(merged) > cap: - merged = self._summarize_batch([{"role": "user", "content": merged}], "") or merged + merged = self._summarize_batch([{"role": "user", "content": merged}]) or merged return CompactionResult( summary=merged or previous_summary, condensed=condensed, stranded=stranded ) - def _summarize_batch(self, batch: list[ChatMessage], previous_summary: str) -> str: - """Fold one batch of dropped turns into *previous_summary*. + def _summarize_batch(self, batch: list[ChatMessage]) -> str: + """Fold one batch of dropped turns into notes. + + Each batch is summarized on its own, with no carried-forward notes in + the prompt: summarize_history merges the per-batch notes instead, which + keeps summary depth at one rather than re-summarizing the summary once + per batch. An overflowing batch splits in half and each half folds on its own: batch sizing is estimate-based, and the cost of an estimate miss here is stranded turns, not a slow call. Depth is log2 of the batch. - Returns *previous_summary* unchanged on any other failure: older notes - beat dropping the turns on the floor. + Returns "" on any other failure, so the caller counts the batch as + stranded rather than reporting turns it has no notes for. """ transcript = "\n".join(f"{m['role']}: {m['content']}" for m in batch) - previous = f"Earlier notes:\n{previous_summary}\n\n" if previous_summary.strip() else "" prompt = COMPACT_PROMPT.format( words=summary_word_budget(self._config.chat_n_ctx_target), - previous=previous, transcript=transcript, ) try: @@ -764,23 +768,23 @@ def _summarize_batch(self, batch: list[ChatMessage], previous_summary: str) -> s reasoning = split_reasoning(response.text).reasoning.strip() if reasoning: return reasoning - log.warning("History compaction returned nothing; keeping the previous summary") + log.warning("History compaction returned nothing for this batch") except ProviderError as exc: # A single message too big for the window cannot split; it falls # through to the warning below. if exc.kind is ProviderErrorKind.CONTEXT_OVERFLOW and len(batch) > 1: mid = len(batch) // 2 - first = self._summarize_batch(batch[:mid], previous_summary) - second = self._summarize_batch(batch[mid:], "") + first = self._summarize_batch(batch[:mid]) + second = self._summarize_batch(batch[mid:]) merged = "\n".join(part for part in (first, second) if part.strip()) if merged.strip(): return merged - log.warning("History compaction failed; keeping the previous summary", exc_info=True) + log.warning("History compaction failed for this batch", exc_info=True) except Exception: # warning, not debug: the user is told turns were dropped, so the # reason must be in the log by default. - log.warning("History compaction failed; keeping the previous summary", exc_info=True) - return previous_summary + log.warning("History compaction failed for this batch", exc_info=True) + return "" def _known_item_results(self, question: str) -> list[SearchChunk]: """Resolve a document named in *question* to its own chunks. diff --git a/tests/test_compaction.py b/tests/test_compaction.py index 9db44b6d2..d6903136e 100644 --- a/tests/test_compaction.py +++ b/tests/test_compaction.py @@ -89,7 +89,7 @@ def test_batches_each_fit_the_current_model_window() -> None: """ dropped = _msgs(200) # ~20k tokens, i.e. a large conversation ctx = 2048 - batches = batch_overflow(dropped, "", ctx_target=ctx) + batches = batch_overflow(dropped, ctx_target=ctx) assert len(batches) > 1, "a 20k backlog must not be summarized in one 2k call" # `estimate < ctx` is the invariant that shipped a live failure: a batch # estimated at 1728 tokens reached a 2048-token server as ~2666 real tokens @@ -108,25 +108,17 @@ def test_batches_each_fit_the_current_model_window() -> None: assert sum(len(b) for b in batches) == len(dropped) -def test_batches_leave_room_for_the_previous_summary() -> None: - """The previous notes ride in every batch prompt, so they must be budgeted.""" - dropped = _msgs(60) - lean = batch_overflow(dropped, "", ctx_target=4096) - fat = batch_overflow(dropped, "s" * 6000, ctx_target=4096) - assert len(fat) >= len(lean), "a long previous summary must shrink the batches" - - def test_a_single_oversized_turn_is_clipped_rather_than_sent_to_fail() -> None: """One turn bigger than the window would fail every call and lose the lot.""" huge = [{"role": "user", "content": "y" * 200_000}] - batches = batch_overflow(huge, "", ctx_target=2048) + batches = batch_overflow(huge, ctx_target=2048) assert len(batches) == 1 assert estimate_tokens(batches[0][0]) < 2048 assert batches[0][0]["content"].endswith("[…clipped]") def test_no_overflow_yields_no_batches() -> None: - assert batch_overflow([], "", ctx_target=2048) == [] + assert batch_overflow([], ctx_target=2048) == [] def test_an_oversized_turn_flushes_the_batch_being_built() -> None: @@ -137,7 +129,7 @@ def test_an_oversized_turn_flushes_the_batch_being_built() -> None: {"role": "user", "content": "y" * 200_000}, # bigger than any batch {"role": "assistant", "content": "small three"}, ] - batches = batch_overflow(dropped, "", ctx_target=2048) + batches = batch_overflow(dropped, ctx_target=2048) # the two small turns batch together, the giant one stands alone (clipped) assert [len(b) for b in batches] == [2, 1, 1] assert batches[0][0]["content"] == "small one" diff --git a/tests/test_entities_ingest.py b/tests/test_entities_ingest.py index b13408e9d..d445fd58e 100644 --- a/tests/test_entities_ingest.py +++ b/tests/test_entities_ingest.py @@ -112,7 +112,7 @@ def test_spacy_types_load_the_model_when_available(self, mock_svc): fake_nlp = mock.MagicMock(return_value=mock.MagicMock(ents=[])) with ( mock.patch("lilbee.retrieval.concepts.concepts_available", return_value=True), - mock.patch("lilbee.retrieval.concepts.nlp._ensure_spacy_model", return_value=fake_nlp), + mock.patch("lilbee.retrieval.concepts.nlp.load_spacy_pipeline", return_value=fake_nlp), ): rows = asyncio.run(_build_entity_records(_records(), "a.txt")) assert rows is not None @@ -124,7 +124,7 @@ def test_spacy_model_import_error_degrades(self, mock_svc, caplog): with ( mock.patch("lilbee.retrieval.concepts.concepts_available", return_value=True), mock.patch( - "lilbee.retrieval.concepts.nlp._ensure_spacy_model", + "lilbee.retrieval.concepts.nlp.load_spacy_pipeline", side_effect=ImportError("no model"), ), caplog.at_level("WARNING"), diff --git a/tests/test_entities_lifecycle.py b/tests/test_entities_lifecycle.py index d50a4175d..eb9f0325e 100644 --- a/tests/test_entities_lifecycle.py +++ b/tests/test_entities_lifecycle.py @@ -210,7 +210,7 @@ def test_spacy_and_llm_kinds_wire_their_tools(self, isolated): fake_nlp = mock.MagicMock(return_value=mock.MagicMock(ents=[])) with ( mock.patch("lilbee.retrieval.concepts.concepts_available", return_value=True), - mock.patch("lilbee.retrieval.concepts.nlp._ensure_spacy_model", return_value=fake_nlp), + mock.patch("lilbee.retrieval.concepts.nlp.load_spacy_pipeline", return_value=fake_nlp), ): ensure_entities() fake_nlp.assert_called() @@ -234,7 +234,7 @@ def test_spacy_model_import_error_degrades(self, isolated, caplog): with ( mock.patch("lilbee.retrieval.concepts.concepts_available", return_value=True), mock.patch( - "lilbee.retrieval.concepts.nlp._ensure_spacy_model", + "lilbee.retrieval.concepts.nlp.load_spacy_pipeline", side_effect=ImportError("no model"), ), caplog.at_level("WARNING"), diff --git a/tests/test_query.py b/tests/test_query.py index cbb38a11c..5f5cdd9aa 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -1707,6 +1707,17 @@ def test_passage_is_embedded_as_a_query(self, mock_svc): get_services().searcher._hyde_search("explain X", top_k=5) mock_svc.embedder.embed_query.assert_called_once_with("hypothetical passage") + def test_spends_its_own_token_budget(self, mock_svc): + """A generated answer passage and a list of query variants are not the + same length of output, so they must not share one cap: retuning either + one through a shared constant silently retunes the other.""" + from lilbee.retrieval.query.expansion import HYDE_MAX_TOKENS + + mock_svc.provider.chat.return_value = _text_result("hypothetical passage") + mock_svc.store.search.return_value = [] + get_services().searcher._hyde_search("explain X", top_k=5) + assert mock_svc.provider.chat.call_args.kwargs["options"]["num_predict"] == HYDE_MAX_TOKENS + def test_returns_empty_on_error(self, mock_svc): mock_svc.provider.chat.side_effect = RuntimeError("fail") assert get_services().searcher._hyde_search("test", top_k=5) == [] diff --git a/tests/test_summarize_history.py b/tests/test_summarize_history.py index dc8bf3fd0..278c74e8f 100644 --- a/tests/test_summarize_history.py +++ b/tests/test_summarize_history.py @@ -197,7 +197,7 @@ def test_merged_notes_over_the_cap_get_one_compression_pass() -> None: long_note = "note " * 300 # ~375 tokens, well over summary_cap(2048) provider = _provider(long_note) cfg.chat_n_ctx_target = 2048 - plan = plan_compaction(_msgs(60), "", ctx_target=2048) + plan = plan_compaction(_msgs(60), ctx_target=2048) result = _searcher(provider).summarize_history(_msgs(60)) assert provider.chat.call_count == len(plan.batches) + 1, "one merge pass on top of the batches" assert result.summary From 669acc8dfdecf5bc510959836031b5b5b90495f5 Mon Sep 17 00:00:00 2001 From: Tobias <5562156+tobocop2@users.noreply.github.com> Date: Mon, 20 Jul 2026 14:23:34 -0400 Subject: [PATCH 28/28] Fix two servers being able to run against the same data directory (#568) * Canonicalize the data root so one directory keys one lock The server session file, the port file, and the store's write lock are all keyed on paths derived from the data root, spelled exactly as it arrived. Two spellings of the same directory therefore key two different locks and let two servers run against the same data. A symlinked root, a relative root, a leading ~ that systemd units and .env files do not expand, and macOS's /var vs /private/var all produce such a pair. canonical_data_root() resolves a root to its one true path, and every entry point that accepts a raw one now goes through it: the environment and the local walk-up in Config._resolve_defaults, --data-dir and --global in the CLI (which bypasses the config validator entirely and derived its own child paths), the MCP init tool's vault switch, and the Lilbee library constructor, which resolved but never expanded ~. Canonicalizing at the boundary is what makes the invariant structural: every derived path and the exported LILBEE_DATA carry one spelling, so the config and lock layers agree by construction instead of by both remembering to normalize. The lock code itself is unchanged. Users with a symlinked or relative data root will see the resolved path where the raw one used to appear (on macOS, /private/var for /var). Same folder, one spelling. * Tighten the canonicalization docstrings; expand ~ test covers Windows Windows expanduser() reads USERPROFILE rather than HOME, so the library tilde test set only half of what it needed to on that platform. * Canonicalize the config.toml lookup on both branches, not just the env one The local walk-up and platform-default branches reached the same file through the symlink anyway, but routing all three through one call keeps the toml directory and the resolved root visibly the same value. * Read LILBEE_DATA identically in both places, and stop a test patch leaking into config The root resolution stripped the env var and the config.toml lookup did not, so a padded value sent the data root and its settings file to two different directories, which is the same spelling divergence this branch exists to close. Also narrows the autocomplete OSError test. It patched pathlib.Path.expanduser globally, which the config validator also calls on every field assignment, so the cfg-restoring fixture's teardown raised through it. That error predates this branch (it arrived with the data-root expanduser fix) and turns the base red; patching only the autocomplete module's Path tests the same behavior without reaching the validator. * Scope two Path.resolve test patches to the path under test Canonicalizing the data root means the config validator now resolves it, and because Config validates on assignment, that runs whenever any cfg field is set. Two tests patched Path.resolve to raise unconditionally, so the cfg-restoring fixture's teardown raised through them. Both now raise only for their own path and delegate the rest to the real resolve, matching the pattern the Vulkan ICD test already uses. Same branch covered, no collateral damage to unrelated code that resolves a path. * Canonicalize on strings so a patched os.name cannot break the validator Windows CI failed in the cfg-restoring fixture's teardown with 'cannot instantiate PosixPath on your system', raised from inside canonical_data_root's resolve(). Path() picks its flavour from os.name and then calls object.__new__, which skips PosixPath's own Windows guard. So while a test patches os.name to 'posix' -- two do, to cover the os.name == 'posix' branches in devices.py and hermes.py -- any Path built on Windows is a real PosixPath. expanduser() never noticed because it returns self when there is no leading ~, but resolve() rebuilds itself through type(self)(...), and that hits the guard. Expansion and resolution now run as string operations through os.path, with a single Path built at the end, so nothing reconstructs a foreign-flavour path. Same semantics: os.path.realpath is what resolve(strict=False) is built on. * Trim comments to the point --- src/lilbee/api.py | 3 +- src/lilbee/cli/app.py | 7 ++++ src/lilbee/core/config/model.py | 42 ++++++++++------------- src/lilbee/core/system.py | 15 +++++++++ src/lilbee/mcp_server.py | 6 ++-- tests/test_api.py | 10 ++++++ tests/test_config.py | 59 +++++++++++++++++++++++++++------ tests/test_query.py | 10 ++++-- tests/test_system.py | 13 ++++++-- tests/test_tui_widgets.py | 8 ++--- 10 files changed, 124 insertions(+), 49 deletions(-) diff --git a/src/lilbee/api.py b/src/lilbee/api.py index 366bbc5a7..116f98478 100644 --- a/src/lilbee/api.py +++ b/src/lilbee/api.py @@ -35,6 +35,7 @@ from lilbee.app.ingest import copy_files from lilbee.app.services import build_services, services_scope from lilbee.core.config import Config, cfg, config_scope +from lilbee.core.system import canonical_data_root from lilbee.data.store import LOCAL_OWNER, MemoryKind, MemoryRow if TYPE_CHECKING: @@ -81,7 +82,7 @@ def __init__( if config is not None: self._config = config elif documents_dir is not None: - root = Path(documents_dir).resolve() + root = canonical_data_root(documents_dir) self._config = cfg.model_copy( update={ "data_root": root, diff --git a/src/lilbee/cli/app.py b/src/lilbee/cli/app.py index 2fb7d5170..920524738 100644 --- a/src/lilbee/cli/app.py +++ b/src/lilbee/cli/app.py @@ -65,7 +65,14 @@ def _apply_data_root(root: Path) -> None: Exporting the env var keeps spawn-context worker subprocesses on the same data root after their fresh ``import lilbee``. + + Canonicalized here because this bypasses ``Config._resolve_defaults`` and + derives its own children, so a symlinked or relative ``--data-dir`` would + otherwise key a different lock than another process on the same directory. """ + from lilbee.core.system import canonical_data_root + + root = canonical_data_root(root) cfg.data_root = root cfg.documents_dir = root / "documents" cfg.data_dir = root / "data" diff --git a/src/lilbee/core/config/model.py b/src/lilbee/core/config/model.py index 3cf807fc2..c5b7a2a2b 100644 --- a/src/lilbee/core/config/model.py +++ b/src/lilbee/core/config/model.py @@ -44,19 +44,6 @@ _UNSET_PATH = Path() -def _expanded(path: Path) -> Path: - """``path`` with a leading ``~`` expanded, or unchanged if it cannot be. - - ``expanduser()`` raises when the OS cannot resolve a home directory (no - HOME, an unknown ``~user``). Configuration must still load in that case, so - the literal path stands rather than taking the whole process down. - """ - try: - return path.expanduser() - except (OSError, RuntimeError): - return path - - class Config(BaseSettings): """Runtime configuration: one singleton instance, mutated by CLI overrides.""" @@ -1032,7 +1019,12 @@ def _parse_ent_types(cls, v: Any) -> frozenset[str]: @model_validator(mode="before") @classmethod def _resolve_defaults(cls, data: Any) -> Any: - from lilbee.core.system import canonical_models_dir, default_data_dir, find_local_root + from lilbee.core.system import ( + canonical_data_root, + canonical_models_dir, + default_data_dir, + find_local_root, + ) if not isinstance(data, dict): return data @@ -1048,14 +1040,10 @@ def _resolve_defaults(cls, data: Any) -> Any: else: local = find_local_root() data["data_root"] = local if local is not None else default_data_dir() - # data_root may arrive as a raw string (e.g. from LILBEE_DATA_ROOT); the - # child-path derivations below use ``/``, so coerce to Path first. - # expanduser() so a "~/lilbee" value from a systemd unit or .env (which - # do not expand ~) points at the home dir instead of creating a literal - # ./~ tree that a server lock keyed on this path would then diverge on. - # It raises when the OS cannot resolve a home directory, and config must - # still load then, so an unresolvable ~ keeps the literal path. - root = _expanded(Path(data["data_root"])) + # Every child path below derives from this, and the server lock keys on + # those, so canonicalizing here is what makes one directory key one lock. + # Also coerces a raw string (LILBEE_DATA_ROOT) to Path. + root = canonical_data_root(data["data_root"]) data["data_root"] = root if data.get("documents_dir") in (None, _UNSET_PATH): data["documents_dir"] = root / "documents" @@ -1077,15 +1065,19 @@ def settings_customise_sources( dotenv_settings: Any, file_secret_settings: Any, ) -> tuple[Any, ...]: - from lilbee.core.system import default_data_dir, find_local_root + from lilbee.core.system import canonical_data_root, default_data_dir, find_local_root - data_env = os.environ.get("LILBEE_DATA", "") + # .strip() to match _resolve_defaults; a padded value would otherwise + # send the root and its config.toml to different directories. + data_env = os.environ.get("LILBEE_DATA", "").strip() if data_env: toml_dir = Path(data_env) else: local = find_local_root() toml_dir = local if local else default_data_dir() - toml_path = toml_dir / "config.toml" + # Same call as the root itself, so this looks where the root resolves to; + # a "~/lilbee" value would otherwise search a literal ./~ and find nothing. + toml_path = canonical_data_root(toml_dir) / "config.toml" plain_env = _PlainEnvSource(settings_cls, env_prefix="LILBEE_", env_ignore_empty=True) sources: list[Any] = [init_settings, plain_env] diff --git a/src/lilbee/core/system.py b/src/lilbee/core/system.py index 1a22d3583..a54ebdb86 100644 --- a/src/lilbee/core/system.py +++ b/src/lilbee/core/system.py @@ -66,6 +66,21 @@ def find_local_root(start: Path | None = None) -> Path | None: return None +def canonical_data_root(root: Path | str) -> Path: + """Resolve a data root to one canonical path. + + Session file, port file, and write lock all derive from the data root, so + two spellings of one directory key two locks. Symlinks, relative paths, a + leading ``~``, and macOS ``/var`` vs ``/private/var`` each produce a pair. + A root that does not exist yet resolves to where it will be created. + + Uses ``os.path`` rather than ``Path.expanduser().resolve()``: ``resolve`` + rebuilds via ``type(self)``, which raises for a ``PosixPath`` that exists + on Windows (``Path()`` picks its flavour from ``os.name``, which tests patch). + """ + return Path(os.path.realpath(os.path.expanduser(os.fspath(root)))) + + def canonical_models_dir() -> Path: """Return the shared models directory (always in the platform default, never per-project). Multiple lilbee instances share this directory so models are downloaded once. diff --git a/src/lilbee/mcp_server.py b/src/lilbee/mcp_server.py index 603b04536..657b9e0d6 100644 --- a/src/lilbee/mcp_server.py +++ b/src/lilbee/mcp_server.py @@ -51,7 +51,7 @@ from lilbee.core.config import cfg from lilbee.core.config.enums import CrawlRenderMode from lilbee.core.settings import overlay_persisted_settings -from lilbee.core.system import LOCAL_ROOT_DIRNAME +from lilbee.core.system import LOCAL_ROOT_DIRNAME, canonical_data_root from lilbee.crawler import crawler_available, is_url, require_valid_crawl_url from lilbee.crawler.task import get_task, start_crawl from lilbee.data.store import ( @@ -408,7 +408,9 @@ def init(path: str = "") -> dict[str, Any]: "init is unavailable on the HTTP server: it is bound to one vault and " "shared by every connected client. Start a separate server for another vault." ) - base = Path(path) if path else Path.cwd() + # Canonical so this vault keys the same lock paths a CLI or server process + # would derive for the same directory. + base = canonical_data_root(path) if path else canonical_data_root(Path.cwd()) root = base / LOCAL_ROOT_DIRNAME created = False diff --git a/tests/test_api.py b/tests/test_api.py index 9249ee719..eb5322d09 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -59,6 +59,16 @@ def test_create_with_documents_dir(self, tmp_path): assert bee.config.data_dir.exists() assert "myproject" in str(bee.config.data_root) + def test_documents_dir_tilde_expands_instead_of_literal_dir(self, tmp_path, monkeypatch): + """A "~/vault" root must reach the home directory, not a literal ./~ tree.""" + from lilbee import Lilbee + + # POSIX expanduser reads HOME; the Windows one reads USERPROFILE. + monkeypatch.setenv("HOME", str(tmp_path)) + monkeypatch.setenv("USERPROFILE", str(tmp_path)) + bee = Lilbee("~/tilde_vault") + assert bee.config.data_root == (tmp_path / "tilde_vault").resolve() + def test_create_with_config(self, tmp_path): from lilbee import Lilbee diff --git a/tests/test_config.py b/tests/test_config.py index 3264323f0..ef2f46c7a 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -116,6 +116,41 @@ def test_data_root_expands_user_home(self): assert c.data_root == Path.home() / "lilbee_expanduser_probe" assert c.lancedb_dir == Path.home() / "lilbee_expanduser_probe" / "data" / "lancedb" + def test_symlinked_data_root_keys_the_same_paths_as_its_target(self, tmp_path): + """Two spellings of one directory derive one set of lock paths.""" + real = tmp_path / "real_root" + real.mkdir() + link = tmp_path / "link_root" + link.symlink_to(real) + + with mock.patch.dict(os.environ, {"LILBEE_DATA": str(real)}): + direct = Config() + with mock.patch.dict(os.environ, {"LILBEE_DATA": str(link)}): + through_link = Config() + + assert through_link.data_root == direct.data_root + assert through_link.data_dir == direct.data_dir + assert through_link.lancedb_dir == direct.lancedb_dir + + def test_padded_data_env_finds_the_same_dir_for_root_and_config( + self, tmp_path, overlay_reads_config_toml + ): + """A padded LILBEE_DATA sends the root and its config.toml to one dir.""" + (tmp_path / "config.toml").write_text("top_k = 7\n", encoding="utf-8") + with mock.patch.dict(os.environ, {"LILBEE_DATA": f" {tmp_path} "}): + c = Config() + assert c.data_root == tmp_path + assert c.top_k == 7 + + def test_relative_data_root_resolves_absolute(self, tmp_path, monkeypatch): + """A relative root must not re-key on the process working directory.""" + (tmp_path / "kb").mkdir() + monkeypatch.chdir(tmp_path) + with mock.patch.dict(os.environ, {"LILBEE_DATA": "kb"}): + c = Config() + assert c.data_root.is_absolute() + assert c.data_root == (tmp_path / "kb").resolve() + def test_empty_data_root_falls_back_to_default_not_cwd(self): """An empty LILBEE_DATA_ROOT must resolve to the platform default, not the process cwd (which would make the data dir move with the launcher).""" @@ -135,17 +170,19 @@ def test_empty_data_root_falls_back_to_default_not_cwd(self): assert c2.data_root != Path.cwd() assert str(c2.data_root).endswith("lilbee") - def test_unexpandable_home_keeps_the_literal_path(self): - """expanduser() raises when the OS cannot resolve a home directory (no - HOME, an unknown ~user). The data root must still resolve rather than - taking config construction -- and the whole process -- down.""" - from lilbee.core.config.model import _expanded - - literal = Path("~/lilbee") - with mock.patch.object(Path, "expanduser", side_effect=OSError("no home")): - assert _expanded(literal) == literal - # The normal case still expands. - assert _expanded(literal) == Path.home() / "lilbee" + def test_unresolvable_home_still_loads_config(self): + """An unresolvable ~ still yields a usable root instead of raising. + + os.path.expanduser returns an unknown ~user unchanged; Path.expanduser + raises. + """ + from lilbee.core.system import canonical_data_root + + root = canonical_data_root("~nosuchuser_lilbee_probe/lilbee") + assert root.is_absolute() + assert str(root).endswith("lilbee") + # The normal case still expands to the real home. + assert canonical_data_root("~/lilbee") == Path.home() / "lilbee" def test_local_server_urls_from_env(self, tmp_path): env = _clean_env(tmp_path) diff --git a/tests/test_query.py b/tests/test_query.py index 5f5cdd9aa..e29c9dd44 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -176,10 +176,16 @@ def test_falls_back_to_raw_on_resolve_failure(self, tmp_path, monkeypatch): from lilbee.retrieval.query import display_source_path cfg.documents_dir = tmp_path / "docs" + target = cfg.documents_dir / "anything.md" + + # Raise only for the path under test; an unconditional patch also hits + # the config validator and blows up in fixture teardown. + real_resolve = _Path.resolve - # Force resolve() to raise so the fallback path runs. def _raise(self, strict=False): - raise OSError("simulated") + if self == target: + raise OSError("simulated") + return real_resolve(self, strict=strict) monkeypatch.setattr(_Path, "resolve", _raise) assert display_source_path("anything.md") == "anything.md" diff --git a/tests/test_system.py b/tests/test_system.py index 9f82cedef..35f94e6b4 100644 --- a/tests/test_system.py +++ b/tests/test_system.py @@ -67,11 +67,18 @@ def test_is_network_path_false_when_mounts_unreadable(self, tmp_path, monkeypatc @pytest.mark.skipif(sys.platform == "win32", reason="POSIX mount-path semantics") def test_is_network_path_uses_raw_path_when_resolve_fails(self, mounts_file, monkeypatch): # Path.resolve has no injectable seam, so this one branch patches it. - def _raise(self): - raise OSError("resolve failed") + # Raise only for the path under test; an unconditional patch also hits + # the config validator and blows up in fixture teardown. + target = Path("/workspace/models/m.gguf") + real_resolve = Path.resolve + + def _raise(self, strict=False): + if self == target: + raise OSError("resolve failed") + return real_resolve(self, strict=strict) monkeypatch.setattr(Path, "resolve", _raise) - assert is_network_path(Path("/workspace/models/m.gguf")) is True + assert is_network_path(target) is True class TestHelpers: diff --git a/tests/test_tui_widgets.py b/tests/test_tui_widgets.py index c99a40419..7ef07ec05 100644 --- a/tests/test_tui_widgets.py +++ b/tests/test_tui_widgets.py @@ -2755,12 +2755,10 @@ def test_add_arg_completions(self, tmp_path: object) -> None: assert any("testfile.txt" in x for x in r) def test_path_exists_swallows_oserror(self, monkeypatch: pytest.MonkeyPatch) -> None: - """A path the OS refuses to stat must read as missing, not raise. + """A path the OS refuses to stat reads as missing, not raise. - Patches this module's ``Path`` rather than ``pathlib.Path.expanduser`` - itself: a global patch also breaks the config validator, which expands - the data root on every field assignment, and so blows up in the - cfg-restoring fixture's teardown before monkeypatch unwinds. + Patches this module's Path, not pathlib.Path.expanduser: a global patch + also hits the config validator and blows up in fixture teardown. """ from lilbee.cli.tui.widgets import autocomplete