diff --git a/docs/architecture.md b/docs/architecture.md index fce977469..efe0934d7 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. 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**: 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. 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/docs/usage.md b/docs/usage.md index 1a5fb3f69..6411714dc 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -814,10 +814,17 @@ 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. 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 | +| `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 | | `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/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/app/settings_map.py b/src/lilbee/app/settings_map.py index eaba796e6..11eb1d6d4 100644 --- a/src/lilbee/app/settings_map.py +++ b/src/lilbee/app/settings_map.py @@ -857,6 +857,42 @@ 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)", + ), + "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)", + ), + "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", + ), + "filter_structural_chunks": SettingDef( + bool, + nullable=False, + group=SettingGroup.RETRIEVAL, + help_text="Drop tables-of-contents and classification-banner cover pages from results", + ), "history_rewrite": SettingDef( bool, nullable=False, @@ -903,6 +939,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/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/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/core/config/model.py b/src/lilbee/core/config/model.py index a12647cfa..c5b7a2a2b 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 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( @@ -194,6 +195,42 @@ 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) + + # Lexical (BM25) arm weight relative to the vector arm in rank fusion. + # 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. + lexical_fusion_weight: float = ConfigField(default=1.0, ge=0.0, le=1.0, writable=True) + + # 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 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 # 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. @@ -243,6 +280,15 @@ 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. + # 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) @@ -973,11 +1019,20 @@ 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 + # 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: @@ -985,7 +1040,11 @@ 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"] + # 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" if data.get("data_dir") in (None, _UNSET_PATH): @@ -1006,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/data/export.py b/src/lilbee/data/export.py index 3ad1e080c..3754ee815 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: @@ -94,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 @@ -200,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], @@ -225,6 +275,13 @@ 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) + # 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 # 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. @@ -236,8 +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=meta, ) ], ) diff --git a/src/lilbee/data/ingest/extract.py b/src/lilbee/data/ingest/extract.py index 0aea58d4f..4474364e4 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 derive_title, 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 ( @@ -507,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, @@ -514,23 +531,30 @@ 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. """ 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 [] + 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) - 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( @@ -538,7 +562,7 @@ async def _handle_image( ) if not chunks: _warn_empty_ocr(source_name, "images") - return chunks + return chunks, meta async def ingest_document( @@ -549,11 +573,12 @@ async def ingest_document( quiet: bool = False, on_progress: DetailedProgressCallback = noop_callback, page_texts_out: list[PageTextRecord] | 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 ``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. @@ -567,8 +592,12 @@ 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. + 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, @@ -577,9 +606,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) @@ -599,7 +629,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, @@ -614,6 +644,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("# "): + return stripped[2:].strip() or None + return None + return None async def ingest_markdown( @@ -621,15 +668,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. @@ -637,7 +687,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")) @@ -645,7 +695,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", @@ -660,3 +710,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 1173cb6ee..b21419b85 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, @@ -141,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 @@ -198,22 +200,27 @@ async def _produce_records( quiet: bool = False, on_progress: DetailedProgressCallback = noop_callback, page_texts_out: list[PageTextRecord] | 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 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 [] 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: - records = await ingest_document( + records, meta = await ingest_document( path, source_name, content_type, @@ -222,7 +229,11 @@ async def _produce_records( page_texts_out=page_texts, ) - return records + 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 + return records, meta def _disk_stat(path: Path) -> SourceStat | None: @@ -524,6 +535,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() @@ -680,7 +692,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] = [] - records = await _produce_records( + records, meta = await _produce_records( entry.path, name, entry.content_type, @@ -707,6 +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, ) except (asyncio.CancelledError, TaskCancelledError) as exc: # TaskCancelledError is the TUI's cooperative cancel signal raised @@ -1031,6 +1044,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..af34b16ed --- /dev/null +++ b/src/lilbee/data/ingest/title.py @@ -0,0 +1,46 @@ +"""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 isinstance(metadata_title, str) 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. Extractor metadata is + untyped, so a string ``authors`` is treated as one author (not split into its + characters) and non-string entries are coerced. + """ + 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), + 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..83f02278c 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); None + # when the title is empty, so chunk rows persist NULL like the _sources table. + title: NotRequired[str | None] 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..e210a933e 100644 --- a/src/lilbee/data/store/core.py +++ b/src/lilbee/data/store/core.py @@ -4,14 +4,13 @@ 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 from typing import TYPE_CHECKING import pyarrow as pa -import pyarrow.compute as pc from lilbee.core.config import ( CHUNK_CONCEPTS_TABLE, @@ -28,10 +27,12 @@ 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_weight_scale, fuse_arms, normalized_bm25, vector_similarity from .lance_helpers import ( + _CHUNK_COLUMN, _chunk_type_predicate, _has_fts_index, + _has_scalar_index, _has_vector_index, _safe_delete_unlocked, _sources_search_filter, @@ -64,6 +65,7 @@ PageTextRecord, RemoveResult, SearchChunk, + SourceMeta, SourceRecord, SourceStat, SourceStatBackfill, @@ -105,6 +107,41 @@ def _drop_unsupported_far_rows( _MAX_THRESHOLD = 1.0 _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, + limit: int, + chunk_type: ChunkType | None, + column: str = _CHUNK_COLUMN, +) -> 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. + """ + 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; # refine_factor re-ranks the PQ candidates against full vectors to recover recall. _VECTOR_METRIC = "cosine" @@ -118,6 +155,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. @@ -137,6 +182,12 @@ def _drop_unsupported_far_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.""" @@ -174,6 +225,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. @@ -218,10 +273,25 @@ 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)"}) + # 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: """Return the persisted store metadata row, or ``None`` if unset.""" table = self.open_table(META_TABLE) @@ -410,15 +480,131 @@ def ensure_fts_index(self) -> None: return try: if _has_fts_index(table): - table.optimize() - log.debug("FTS index optimized on '%s'", CHUNKS_TABLE) + # 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 + # table, the title index included. + table.optimize() + log.debug("FTS index optimized on '%s'", CHUNKS_TABLE) + 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: - table.create_fts_index("chunk", replace=False) + # 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_COLUMN, replace=False, with_position=False) + self._fts_ready = True log.debug("FTS index created on '%s'", CHUNKS_TABLE) - self._fts_ready = True + # 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) + 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: + # 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: + # 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. + + 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_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") + 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. + + ``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(): + 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: + 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.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. @@ -470,11 +656,11 @@ 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) - 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,11 +683,7 @@ def bm25_probe( if not self._fts_ready: return [] try: - query = table.search(query_text, query_type="fts") - 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) @@ -538,6 +720,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() @@ -600,11 +787,50 @@ def _fts_arm( limit: int, chunk_type: ChunkType | None, ) -> list[SearchChunk]: - """BM25-arm candidates.""" - query = table.search(query_text, 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()] + """BM25-arm candidates over the chunk text.""" + return _lexical_rows(table, query_text, limit, chunk_type) + + def _title_arm( + self, + table: lancedb.table.Table, + query_text: str, + limit: int, + chunk_type: ChunkType | None, + ) -> list[SearchChunk]: + """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 + 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 [] + try: + 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, @@ -615,23 +841,53 @@ 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. + 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: + 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 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 + # 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) + ) 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=lexical_weight, + title_weight=title_weight, + weight_total=weight_total, ) fused = _drop_unsupported_far_rows(fused, max_distance) return fused[:top_k] @@ -841,24 +1097,40 @@ 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. 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 [] + 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) + 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: """Delete the sources' chunks, page texts, and chunk-concept rows. @@ -964,8 +1236,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 +1253,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 +1290,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() @@ -1062,12 +1346,13 @@ 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() 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 +1371,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 +1378,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 ] @@ -1418,6 +1704,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. @@ -1428,6 +1715,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/src/lilbee/data/store/fusion.py b/src/lilbee/data/store/fusion.py index 157784f47..743631f70 100644 --- a/src/lilbee/data/store/fusion.py +++ b/src/lilbee/data/store/fusion.py @@ -1,12 +1,22 @@ -"""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 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. + +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 @@ -22,17 +32,56 @@ 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_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 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 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 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 1.0 + margin = max(0.0, sims[0] - fmean(sims[1:])) + confidence = min(1.0, margin / margin_scale) + 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]: """Scale raw BM25 scores against the list maximum, into (0, 1]. @@ -55,28 +104,64 @@ def _rank_weight(rank: int) -> float: return (_RRF_K + 1) / (_RRF_K + rank) -def fuse_arms( - vector_rows: list[SearchChunk], - fts_rows: list[SearchChunk], -) -> list[SearchChunk]: - """Merge the two arms into one list scored by reciprocal 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. - 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)``. + 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. """ - 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): + for rank, row in enumerate(rows, start=1): key = _key(row) - lexical = _rank_weight(rank) / 2 + contribution = _rank_weight(rank) * share seen = merged.get(key) if seen is None: - merged[key] = row.model_copy(update={"score": lexical}) + merged[key] = row.model_copy(update={"score": contribution}) else: - merged[key] = seen.model_copy( - update={"score": (seen.score or 0.0) + lexical, "bm25_score": row.bm25_score} - ) + 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, + *, + 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. + + 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)``. + + *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). + """ + 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 / 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 / weight_total) + if title_rows and title_weight > 0: + _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/data/store/lance_helpers.py b/src/lilbee/data/store/lance_helpers.py index 6565f3ac1..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,17 +133,30 @@ 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_COLUMN) -> 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 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/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..8f0ab2a61 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 @@ -237,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/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/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/dedup.py b/src/lilbee/retrieval/query/dedup.py index a1159a400..9f5314d99 100644 --- a/src/lilbee/retrieval/query/dedup.py +++ b/src/lilbee/retrieval/query/dedup.py @@ -105,11 +105,11 @@ 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, 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/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/neighbors.py b/src/lilbee/retrieval/query/neighbors.py new file mode 100644 index 000000000..d456a083c --- /dev/null +++ b/src/lilbee/retrieval/query/neighbors.py @@ -0,0 +1,191 @@ +"""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. + + 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. + """ + 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). + """ + 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) + 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. + + 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(): + wanted = sorted( + { + index + for center in owned + for index in range(center - radius, center + radius + 1) + if index >= 0 + } + ) + 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 + 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: + 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 1d293f31c..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, @@ -75,6 +76,8 @@ title_candidates, ) 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, @@ -225,10 +228,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: @@ -423,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. @@ -625,6 +634,19 @@ 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 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: + # 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) + 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: @@ -679,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) @@ -697,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: @@ -743,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. @@ -916,33 +941,31 @@ def _finalize_context( retrieved set tighter without re-running retrieval or condensation. """ system = self._system_with_memory(self._config.rag_system_prompt, question) - results = self._fit_context_budget(results, system, question, history, scale) + base_results = list(results) + 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}] 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: """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 @@ -962,7 +985,21 @@ def _fit_context_budget( + sum(self._budget_tokens(m["content"]) for m in history or []) + _CONTEXT_TEMPLATE_TOKENS ) - budget = int((prompt_token_budget(ctx) - non_source) * scale) + return int((prompt_token_budget(ctx) - non_source) * scale) + + 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. + """ kept: list[SearchChunk] = [] used = 0 for r in results: @@ -977,7 +1014,21 @@ def _fit_context_budget( len(kept), len(results), ) - return kept + return kept, used + + def _widen_with_neighbors(self, results: list[SearchChunk], leftover: int) -> list[SearchChunk]: + """Widen each fitted passage with adjacent same-source chunks. + + 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 or leftover <= 0: + return results + 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.""" @@ -1258,7 +1309,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( @@ -1267,13 +1318,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 ) @@ -1347,10 +1403,10 @@ 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. + # 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/src/lilbee/retrieval/query/structural.py b/src/lilbee/retrieval/query/structural.py new file mode 100644 index 000000000..ed59da2db --- /dev/null +++ b/src/lilbee/retrieval/query/structural.py @@ -0,0 +1,67 @@ +"""Detect document-structure chunks that dilute retrieval precision. + +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 + +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 +# 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. +# 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 + + +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_HITS 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/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_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/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_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_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_config.py b/tests/test_config.py index 7c391100c..ef2f46c7a 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -104,6 +104,86 @@ 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_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).""" + 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") + # 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_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) env["LILBEE_OLLAMA_BASE_URL"] = "http://box:11434" @@ -185,6 +265,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() @@ -747,6 +845,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: 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_export.py b/tests/test_export.py index 2b9259540..4be5f1bbb 100644 --- a/tests/test_export.py +++ b/tests/test_export.py @@ -263,3 +263,47 @@ 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): + # 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) 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..816b8a29b 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 @@ -895,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): @@ -954,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 @@ -1040,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 @@ -1090,13 +1094,21 @@ 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, ): + 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 @@ -2187,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", @@ -2237,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. @@ -2253,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 == [] @@ -2271,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) @@ -2285,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 @@ -2302,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) @@ -2373,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() @@ -2381,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 @@ -2392,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 = "" @@ -2406,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" @@ -2438,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 @@ -2480,7 +2519,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 @@ -2496,7 +2536,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] @@ -2666,7 +2706,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 @@ -2689,7 +2729,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 @@ -2711,7 +2751,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 @@ -2747,7 +2787,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 @@ -2773,7 +2813,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 @@ -2794,7 +2834,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 @@ -2859,7 +2899,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): @@ -2868,7 +2908,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): @@ -2876,7 +2916,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" @@ -2947,7 +2987,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): @@ -2956,7 +2996,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 == [] @@ -3184,3 +3224,61 @@ 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_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" diff --git a/tests/test_ingest_title.py b/tests/test_ingest_title.py new file mode 100644 index 000000000..b29182eed --- /dev/null +++ b/tests/test_ingest_title.py @@ -0,0 +1,72 @@ +"""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 == "" + + 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_neighbors.py b/tests/test_neighbors.py new file mode 100644 index 000000000..1613b4a9d --- /dev/null +++ b/tests/test_neighbors.py @@ -0,0 +1,197 @@ +"""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_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_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) + 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, 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) + 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, 2, 3, 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")] + 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. + 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", [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, + # 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 b1a576f6f..e29c9dd44 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, @@ -134,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.""" @@ -163,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" @@ -501,6 +520,59 @@ 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"] + _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.""" + 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_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.""" + cfg.top_k = 1 # window is top_k*2 = 2 + cfg.filter_structural_chunks = True + 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 + + 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.""" + 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.""" + 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: 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.") + 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 on vector distance alone would re-bury exactly the identifier hits @@ -827,7 +899,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 @@ -845,7 +917,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 = ( @@ -862,12 +934,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): @@ -877,18 +949,145 @@ 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 +class TestNeighborExpansion: + """Selected chunks widen with adjacent same-source chunks at answer time.""" + + 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 + + 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 + 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 + 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 + 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 {prompt_token_budget(ctx)}" + ) + + class TestAskRaw: def test_returns_structured_result(self, mock_svc): mock_svc.store.search.return_value = [_make_result(chunk="oil is 5 quarts")] @@ -1514,6 +1713,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) == [] @@ -2052,7 +2262,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() @@ -2099,7 +2309,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. @@ -2119,7 +2329,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 @@ -2159,7 +2369,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): @@ -2430,7 +2640,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"] @@ -2834,7 +3044,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): @@ -3375,7 +3585,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_settings.py b/tests/test_settings.py index 20f167557..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,12 +209,37 @@ 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): + + 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 + + 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.""" + + 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).""" @@ -242,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 @@ -251,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 @@ -264,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 @@ -273,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"] @@ -281,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 @@ -289,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 @@ -299,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 @@ -309,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"] @@ -318,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 @@ -457,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. @@ -496,3 +510,28 @@ 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): + + 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): + + 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): + + 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..642938629 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -11,6 +11,7 @@ CitationRecord, SearchChunk, SearchScope, + SourceMeta, SourceType, Store, cosine_sim, @@ -162,8 +163,107 @@ def test_second_call_optimizes_instead_of_rebuilding(self, store): create_spy.assert_not_called() 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.""" + 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_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, caplog): + """An unrelated optimize() failure keeps the existing index and does NOT + 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") + 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, + 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 + 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.""" store.add_chunks(_make_records()) table = store.open_table("chunks") assert table is not None @@ -171,10 +271,40 @@ 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"] + 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 + + 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", "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) + # Both indexes are positionless: with_position=True overflows LanceDB's + # 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_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, 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() + 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 @@ -197,6 +327,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 @@ -256,6 +395,108 @@ 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 + + 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.""" @@ -602,6 +843,56 @@ 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_feeds_a_derived_weight_to_fusion(self, store, test_config): + """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()) + store.ensure_fts_index() + query_vec = [0.5] * 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.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") + 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 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_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 + test_config.lexical_fusion_weight = 0.3 + store.add_chunks(_make_records()) + store.ensure_fts_index() + query_vec = [0.5] * test_config.embedding_dim + from lilbee.data.store import core as store_core + + with ( + 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") + adapt.assert_not_called() + assert fuse.call_args.kwargs["lexical_weight"] == pytest.approx(0.3) + assert len(results) > 0 + def test_fallback_to_vector_when_no_query_text(self, store, test_config): records = _make_records() store.add_chunks(records) @@ -1191,31 +1482,58 @@ 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) + 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"] - # 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 _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 - def _raise_search(*args, **kwargs): - raise AttributeError("LanceFtsQueryBuilder has no attribute 'metric'") - table.search = _raise_search - return table +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) - 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" + 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_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( @@ -2195,3 +2513,304 @@ 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. 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)) + 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 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 + 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) + # 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_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 + 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 + + 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") + 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_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.""" + 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_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")) + 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, 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") + 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), + 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: + """Extraction-time document metadata persisted on the sources table.""" + + def test_upsert_source_persists_meta(self, store): + 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 + + 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..28035ced8 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): @@ -99,6 +104,126 @@ 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.""" + + 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.""" + + 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_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)} + explicit = {r.source: r.score for r in fuse_arms(vec, lex, lexical_weight=1.0)} + assert default == explicit + + 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 + + class TestRegressionMechanism: """Synthetic reproduction of the graded-A/B regression shape: a lexical query whose relevant passages are mutually similar (they all quote the @@ -186,3 +311,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) diff --git a/tests/test_structural.py b/tests/test_structural.py new file mode 100644 index 000000000..f3d8fb312 --- /dev/null +++ b/tests/test_structural.py @@ -0,0 +1,83 @@ +"""Tests for the document-structure (TOC / cover-page) chunk detector.""" + +from lilbee.retrieval.query.structural import is_structural_chunk + +# A table of contents: section titles followed by dot leaders and page numbers. +TOC = """Contents +A. Executive Summary ................................................. 1 +B. Introduction ..................................................... 3 +C. Regional Overview ............................................... 9 +D. Budget and Staffing ............................................. 10 +E. Program Findings ................................................ 11 +F. Recommendations ................................................. 13 +""" + +# A classification-banner cover/title page. +COVER = """UNCLASSIFIED +NATIONAL PROGRAM REVIEW BOARD +OFFICE OF STRATEGIC ASSESSMENT +Fiscal Year 2024 Consolidated Annual Report on +Program Performance and Oversight +Information Cut Off Date: 1 JUNE 2024 +UNCLASSIFIED +""" + +# Real prose that must NOT be flagged. +PROSE = ( + "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 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." +) + + +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("REGIONAL OVERVIEW AND PROGRAM FINDINGS") is False + + def test_short_classified_body_page_is_not_a_cover(self): + # 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 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 + + 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 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 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 0f141bc7e..7ef07ec05 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,21 @@ 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 reads as missing, not raise. + 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 - 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: