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 1/4] 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 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 2/4] 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 3/4] 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 4/4] 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())